Improved selection of fit functions for Q spectra using AIC to balance fit...

    return np.interp(val, x[:N], y[:N])

# Main solver

def xscomptb(energy,  ktbb, kte, alpha, logA, Norm):

def xscomptb(energy,  ktbb, bb_par, kte, alpha, logA, Norm):

    # Unpack parameters

    delta=0

    bb_par=3

    N=4000

    N=1000

    #ktbb, alpha, kte, logA = parameters

    A = 10**logA

import numpy as np

from scipy.special import gamma as gammln, zeta

from numba import njit

# Funzione blackbody modificata

@njit

def bbody(R, z, bb_par):

    return z**bb_par / (np.exp(R * z) - 1.0)

# Interpolazione lineare custom compatibile con Numba

@njit

def interp_funct(x, y, N, val):

    result = np.zeros_like(val)

    for i in range(len(val)):

        v = val[i]

        if v <= x[0]:

            result[i] = y[0]

        elif v >= x[N-1]:

            result[i] = y[N-1]

        else:

            for j in range(N - 1):

                if x[j] <= v < x[j + 1]:

                    t = (v - x[j]) / (x[j + 1] - x[j])

                    result[i] = y[j] * (1 - t) + y[j + 1] * t

                    break

    return result

# Parte numerica pesante accelerata con Numba

@njit

def xscomptb_numba(energy, ktbb, bb_par, kte, alpha, logA, Norm):

    delta = 0.0

    N = 1000

    A = 10 ** logA

    fb = 1.0

    R = kte / ktbb

    cn = 8.0525 * (kte / ktbb) ** bb_par * ktbb ** (bb_par - 4)

    gamma_val = alpha * (alpha + delta + 3.0) + 3.0 * delta

    tmin = -10.0

    tmax = 10.0

    h = (tmax - tmin) / N

    t = tmin + h * np.arange(N + 1)

    xnum = np.exp(t)

    p = np.full(N + 1, fb)

    g = xnum - 3.0 * fb - delta

    r = xnum - gamma_val + 3.0 * delta

    s = -bbody(R, xnum, bb_par)

    for i in range(N + 1):

        if s[i] > -1e-15:

            s[i] = 0.0

    a = p / h ** 2

    b = -(2 * p / h ** 2 + g / h - r)

    c = p / h ** 2 + g / h

    # Condizioni al contorno

    a[0], b[0], c[0], s[0] = 0.0, 1.0, 0.0, 0.0

    a[N], b[N], c[N], s[N] = 0.0, 1.0, 0.0, 0.0

    # Algoritmo di Thomas (tridiagonale)

    L = np.zeros(N + 1)

    K = np.zeros(N + 1)

    L[0] = -c[0] / b[0]

    K[0] = 0.0

    for n in range(1, N + 1):

        denom = a[n] * L[n - 1] + b[n]

        L[n] = -c[n] / denom

        K[n] = (s[n] - a[n] * K[n - 1]) / denom

    # Back-substitution

    spec = np.zeros(N + 1)

    integ = 0.0

    for n in range(N - 1, -1, -1):

        spec[n] = L[n] * spec[n + 1] + K[n]

        integ += spec[n] * h

    Nflux = len(energy)

    flux_bb = np.zeros(Nflux)

    unnorm_flux = np.zeros(Nflux)

    x = energy / kte

    flux_comp = interp_funct(xnum, spec, N, x) / energy

    bb = bbody(R, x, bb_par) / energy

    for i in range(1, Nflux):

        unnorm_flux[i] = 0.5 * (flux_comp[i] + flux_comp[i - 1])

        flux_bb[i] = 0.5 * (bb[i] + bb[i - 1])

    return flux_bb, unnorm_flux, integ

# Funzione wrapper finale (usa funzioni speciali non compatibili con Numba)

def xscomptb(energy, ktbb, bb_par, kte, alpha, logA, Norm):

    flux_bb, unnorm_flux, integ = xscomptb_numba(energy, ktbb, bb_par, kte, alpha, logA, Norm)

    A = 10 ** logA

    phtot = (ktbb ** bb_par) * gammln(bb_par) * zeta(bb_par) / (integ * kte ** bb_par)

    flux = Norm / (A + 1.0) * (flux_bb + A * phtot * unnorm_flux)

    return flux

import numpy as np

from lmfit import Model

#==============================================================================

def smooth_func(x, x0, y0, alpha, k, idx):

#========================================================================================

    #x0 (dove la curva "piega")

    #y0 (il valore della parte piatta)

    #alpha (negativo se decresce, positivo se cresce)

    #k (la "rapidità" della transizione)

def piecewise_bb_cpl_smooth_deriv(E, A, kTbb, Gamma, Ecut, Esplit):

    """

    Funzione spezzata BB+CPL con continuità di valore e derivata in Esplit

    e una sola normalizzazione libera A.

    """

    # Calcolo dell'argomento dell'esponenziale

    # Usiamo np.divide per gestire in modo sicuro il caso x/x0

    # Corpo nero (non normalizzato)

    BB = (E**2) / (np.exp(E / kTbb) - 1)

    # Punto di raccordo

    x = Esplit

    BB_split = (x**2) / (np.exp(x / kTbb) - 1)

    # Derivata del BB in Esplit

    expx = np.exp(x / kTbb)

    dBB_dx = (

        2 * x / (expx - 1)

        - (x**2) * (expx / (kTbb * (expx - 1)**2))

    )

    with np.errstate(divide='ignore'):

        #arg_log = np.log(x / x0)

        arg_log = x - x0

    # CPL e derivata in Esplit (senza normalizzazione)def fit_log_poly_cpl(energy_center=None, delta_energy=None, array_y=None, emin=None, emax=None, deg=5,

                     split_energy=20, ktbb=None, kte=None, tau=None):

    nan_mask = np.isnan(arg_log)

    #result = minimize_scalar(

    #    lambda split: compute_chisq_for_split(split, energy_center, array_y, emin, deg),

    #    bounds=(4, 30), method='bounded'

    # split_energy=result.x

    #)

    # Verifica se ce ne sono

    if np.any(nan_mask):

        cnt = np.sum(nan_mask)

        print(f"Ci sono {cnt} valori NaN")

    y_err = np.sqrt(np.abs(array_y))

    array_y=array_y/delta_energy

    safe_y_err = np.where(y_err == 0, 1e8, y_err)/delta_energy

    energy_at_max = energy_center[np.argmax(array_y)]

    split_energy=2*energy_at_max

    arg_exp = k * arg_log

    split_energy=10

    value=y0 + alpha * np.log(1 + np.exp(arg_exp)) ** idx

    # ==========================================================

    mask_low = (energy_center > emin) & (energy_center <= split_energy) & (array_y > 0)

    log_x = np.log10(energy_center[mask_low])

    log_y = np.log10(array_y[mask_low])

    return value

    # Fit polinomiale in scala log-log

    poly_coeffs = np.polyfit(log_x, log_y, deg=deg)

    poly_fit = np.poly1d(poly_coeffs)

    # Valore e derivata a split_energy

    log_split = np.log10(split_energy)

    log_y_split = poly_fit(log_split)

#==============================================================================

def const_piecewise_parabola(x, x0, a, b, c, w):

    # Converti in valore reale

    y_split = 10 ** log_y_split

    width=w

    x = np.asarray(x)

    #==========================================================

    # Mask per regione CPL

    # =========================================================

    # Sigmoide per la transizione (smooth step)

    s = 1 / (1 + np.exp(-(x - x0) / width))  # logistic transition

    mask_high = (energy_center > split_energy) & (array_y > 0) & (array_y/safe_y_err > 2) & (energy_center < 40)

    # Componente parabolica centrata in x0

    parabola = a * (x - x0) ** 2 + b * (x - x0)

    while (np.sum(mask_high) < 5):

        split_energy=split_energy-1

        mask_high = (energy_center > split_energy) & (array_y > 0) & (array_y / safe_y_err > 2)

    # Transizione continua tra c e parabola + c

    y = (1 - s) * c + s * (parabola + c)

    return y

#==============================================================================

    x_high = energy_center[mask_high]

    y_high = array_y[mask_high]

def method_lm(x_val, y_val, y_err, fit_config, model_func):

    # Stima iniziale: usa continuità per inizializzare A

    gamma0 = 2

    Ecut0 = split_energy

    model = Model(model_func)

    if kte is not None:

        Ecut0=3*kte

    A0 = y_split / (split_energy ** (-gamma0) * np.exp(-split_energy / Ecut0))

    param_names = fit_config['param_names']

    initial_guess = fit_config['initial_guess']

    bounds = fit_config.get('bounds', {})

    bound_left = bounds.get('left', {})

    bound_right = bounds.get('right', {})

    bounds=([0, 0.5, 0.1], [np.inf, 5, 50])

    params = model.make_params(**{name: initial_guess[name] for name in param_names})

    try:

        with warnings.catch_warnings():

            warnings.simplefilter("ignore", OptimizeWarning)

            popt, _ = curve_fit(

                cutoff_powerlaw,

                x_high,

                y_high,

                p0=[A0, gamma0, Ecut0],

                sigma=y_err[mask_high],

                absolute_sigma=True,

                bounds=bounds,

                maxfev=2000

            )

            if not np.all(np.isfinite(popt)):

                return np.zeros(len(array_y), dtype=bool)

    for name in param_names:

        if name in bound_left:

            params[name].set(min=bound_left[name])

        if name in bound_right:

            params[name].set(max=bound_right[name])

            # Costruisci modello combinato

            result_full = np.zeros_like(energy_center)

    result = model.fit(y_val, params, x=x_val, weights=1 / y_err**2, max_nfev=10000)

            print("A %5.3e gamma %5.3f Ecut %5.4f" % (popt[0], popt[1], popt[2]))

    if not result.success:

        return None

            mask_below=(energy_center < split_energy)

            mask_above = (energy_center > split_energy)

    best_params = [result.params[name].value for name in param_names]

            # Parte polinomiale

            result_full[mask_below] = 10 ** poly_fit(np.log10(energy_center[mask_below]))

    #print(best_params)

            #result_full[mask_low]=np.ones(len(energy_center[mask_low]))

    return best_params

            # Parte CPL

            result_full[mask_above] = cutoff_powerlaw(energy_center[mask_above], *popt)

#==================================================================================

def method_curve_fit(x_val, y_val, y_err, p0, bound_left=None, bound_right=None):

    res = least_squares(

        residuals_const_piecewise_parabola,

        x0=p0,

        args=(x_val, y_val, y_err),

        bounds=(bound_left, bound_right),

        method='trf',

        # loss='linear',  # o 'huber' per robusto agli outlier

        max_nfev=10000

    )

    except RuntimeError:

        # Fit failed (did not converge)

        result_full=np.zeros(len(array_y), dtype=bool)

    print(res.success)

    # Calcolo del chi²

    chi_squared = np.sum(

        ((array_y - result_full) / safe_y_err) ** 2

    )/(len(array_y)-3)

    if not res.success:

        raise RuntimeError("Least Squares Fit did not converge: " + res.message)

    #print(f"Chi² del fit: {chi_squared:.2f}")

    return result_full

    return res

    CPL_split = x**(-Gamma) * np.exp(-x / Ecut)

    dCPL_dx = CPL_split * (-Gamma / x - 1 / Ecut)

    # Fattore di raccordo

    R = BB_split / CPL_split

    R_prime = dBB_dx / dCPL_dx

    # Per derivabilità: imposta R = R_prime (sistema omogeneo)

    # Sostituisci R con media geometrica (più stabile numericamente)

    R = np.sqrt(np.abs(BB_split / CPL_split * dBB_dx / dCPL_dx))

    CPL = R * E**(-Gamma) * np.exp(-E / Ecut)

    # Spezzata

    model = np.where(E <= Esplit, BB, CPL)

    return A * model

#!/usr/bin/env python3

import numpy as np

from model_utilities import *

import os,sys

import random

from scipy.optimize import curve_fit

from comptb import *

from model_utilities import *

#=======================================================================

try:

    from heasp import *

    module_heasp=True

    import _heasp

    print("Import successful")

    #print("Import successful")

except ImportError:

    print("Import failed, continuing...")

    module_heasp = False

TableXspec.setNumIntParams(FreeParams)

TableXspec.setNumAddParams(0)

#============================================================

#Angles with respect to the normal of the slab

#============================================================

n_angles=10

u_bin_edges = np.linspace(0, 1, n_angles+1)

delta_u=u_bin_edges[1]-u_bin_edges[0]

u_centers = 0.5 * (u_bin_edges[:-1] + u_bin_edges[1:])

#u_centers = u_centers[::-1]

#for ii in range(len(u_bin_edges)-1):

#    print('Bin %d [%4.2f : %4.2f]' % (ii, u_bin_edges[ii], u_bin_edges[ii+1]))

#========================================================================

def write_table_value(latmax=None, ktbb=None, kte=None, tau=None, u_obs=None, spin=None, mc_file=None,

                      xspec=None, fitsfile=None, paramvalues=None):

    print('==============================================')

    print('Reading file: %s' % (mc_file))

    C_matrix, Q_matrix = read_mc_simulations(mc_file=mc_file, n_i=20, n_j=256)

                                        delta_energy=delta_energy, ktbb=ktbb, kte=kte, tau=tau)

    #for ii in range(len(ene_center)):

    #    print("[%5.4f - %5.4f]"% (ene_center[ii]-delta_energy[ii]/2, ene_center[ii]+delta_energy[ii]/2))

    PD_filtered = smooth_matrix_pd(Q_matrix=Q_matrix, C_matrix=C_matrix,

                                   energy_center=ene_center, emin=0.3, emax=30)

    #print(C_matrix.shape, Q_matrix.shape, C_filtered.shape, PD_filtered.shape)

    if C_matrix.shape[0] != len(u_centers):

        print('Error: number of u angles from the C_matrix is %d' % (C_matrix.shape[0]))

        print('while the array u_center has %d values ' % (len(u_centers)))

        exit()

    Norm_flux = normalize_counts(counts_array=C_matrix, energy_center=ene_center, delta_energy=delta_energy,

                                 u_array=u_centers, delta_u=delta_u, ktbb=float(ktbb))

    I_matrix = intensity_matrix(counts_array=C_filtered, energy_center=ene_center, delta_energy=delta_energy,

                                delta_u=delta_u, norm_value=Norm_flux)

    return C_filtered, PD_filtered

#============================================================

#Angles with respect to the normal of the slab

#============================================================

u_bin_edges = np.linspace(0, 1, 21)

delta_u=u_bin_edges[1]-u_bin_edges[0]

u_centers = 0.5 * (u_bin_edges[:-1] + u_bin_edges[1:])

#for ii in range(len(u_bin_edges)-1):

#    print('Bin %d [%4.2f : %4.2f]' % (ii, u_bin_edges[ii], u_bin_edges[ii+1]))

#for ii in range(energy_bins-1):

#    print('Channel %d [%5.3f - %5.3f]' % (ii, ene_center[ii]-delta[ii]/2, ene_center[ii]+delta[ii]/2 ))

    if inputval.startswith('ymin') == True:

        ymin = float(inputval[inputval.rfind('=') + 1:])

    if inputval.startswith('ymax') == True:

        ymax = float(inputval[inputval.rfind('=') + 1:])

    if inputval.startswith('tabname') == True:

        tabname = str(inputval[inputval.rfind('=') + 1:])

    if inputval.startswith('timeon') == True:

        timeon = str(inputval[inputval.rfind('=') + 1:])

    if inputval.startswith('emax') == True:

        emax = float(inputval[inputval.rfind('=') + 1:])

try:

    runtable

except:

except:

    ymin=-5

try:

    ymax

except:

    ymax=None

try:

    timeon

except:

    timeon='n'

#====================================================================

if runtable is None:

    ktbb,kte,tau=filename2params(filename=mcfile)

    pdf_name = os.path.basename(mcfile.replace(".dat", ".pdf"))

    pdf_name = pdf_name.replace("test_BL", param)

    plt.figure(figsize=(10, 5))

    for ii in range(1, 2):

    u_val=random.randint(0, len(u_centers)-1)

    #u_val=5

    #for ii in range(len(u_centers)):

    #    print(" %d %4.3f %4.3f" % (ii, u_centers[ii], np.arccos(u_centers[ii])/np.pi*180))

   # print('==========================')

    print("Selected window ", u_val)

    print('Angle %4.3f deg'% (np.arccos(u_centers[u_val])/np.pi*180))

   # print('==========================')

    #u_val=18

    for ii in range(u_val, u_val+1):

        if param == 'Q':

            plot_data(x_value=ene_center, y_mc=Q_matrix[ii], y_smooth=Q_filtered[ii],

            error = np.sqrt(2 / safe_Cmatrix[ii])

            plot_data(x_value=ene_center, y_mc=PD_matrix[ii] * 100, y_smooth=PD_filtered[ii] * 100, ymin=ymin, ymax=5,

                        emin=0.3, emax=30, y_err=error*100, y_label='Polarization degree (%)', pdf_name=pdf_name)

            plot_data(x_value=ene_center, y_mc=PD_matrix[ii] * 100, y_smooth=PD_filtered[ii] * 100, ymin=ymin, ymax=ymax,

                        emin=0.3, emax=emax, y_err=error*100, y_label='Polarization degree (%)', pdf_name=pdf_name,

                      delta_energy=delta_energy, param=param)

        elif param == 'C':

            plot_data(x_value=ene_center, y_mc=C_matrix[ii], y_smooth=C_filtered[ii],

                      logscale=True, ymin=0.1, ymax=1e7, y_err=np.sqrt(C_matrix[ii]),

                      delta_energy=delta_energy, y_label='Counts/keV', pdf_name=pdf_name)

                      logscale=True, ymin=0.1, ymax=1e8, y_err=np.sqrt(C_matrix[ii]),

                      delta_energy=delta_energy, y_label='Counts/keV', pdf_name=pdf_name, param=param, emax=emax)

            #I_flux, Q_flux, U_flux=read_spectra(filename='polarimetry.dat')

            print('Unknown parameter name to be plotted')

            exit(1)

    #plt.pause(2)  # Mostra il plot per 0.5 secondi

    #plt.close()  # Chiude la figura

    #sys.exit()

    FlagTime=False

    if timeon=='y':

        plt.pause(2)  # Mostra il plot per 0.5 secondi

        plt.close()  # Chiude la figura

        sys.exit()

    else:

        plt.show()

        pass

    print('=================================================================================')

mcfile='../grid0/test_BL_slab_Z0seed_grid__2.5_4.0_4.0__IQ_ADMP.dat'

#====================================================================

if runtable=='n':

#kTbb

ktbb_array = np.linspace(0.5,3,6)

ktbb_array = np.arange(1, 3.5, 0.5)  #Last value is excluded by arange

formatted_ktbb_val=[f"{val:.1f}" for val in ktbb_array]

ktbb_init=1.5

ktbb_init=2

table_init(init_val=ktbb_init, array_values=ktbb_array, name='kTbb', table=TableXspec)

#kte_array = np.linspace(1.5,5,8)

kte_array=np.array([1.5,2,2.5,3,3.5,4,4.5,5,6,7])

kte_array=np.array([2,2.5,3,3.5,4,4.5,5,6,7])

formatted_kte_val=[f"{val:.1f}" for val in kte_array]

kte_init=3

tau_array = np.linspace(0.5,4, 8)