Qweak Bayesian Code

import sys
sys.path.append('/home/elhamm/projects/def-wdconinc/elhamm/my_venv/lib/python3.11/site-packages')

import math
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from numpy import random
import arviz as az
import matplotlib as mpl

N_M = 108
N_T = 4
P=0.87
pi=math.pi
A_bias = [0.019]
dA_bias = [0.028]
A_BB = [-0.012]
dA_BB = [0.027]

A_1_M=np.full(N_M,-2.233)
A_2_M=np.full(N_M,-2.233)
A_3_M=np.full(N_M,-3.183)
A_4_M=np.full(N_M,-2.573)
A_5_M=np.full(N_M,-2.093)
A_6_M=np.full(N_M,0.167)
A_7_M=np.full(N_M,1.077)
A_8_M=np.full(N_M,-1.483)

A_1_T=np.full(N_T,2.587)
A_2_T=np.full(N_T,6.147)
A_3_T=np.full(N_T,6.587)
A_4_T=np.full(N_T,2.727)
A_5_T=np.full(N_T,-4.493)
A_6_T=np.full(N_T,-1.193)
A_7_T=np.full(N_T,18.617)
A_8_T=np.full(N_T,-3.993)

A_values_i_M = np.array([-2.233, -2.233, -3.183, -2.573, -2.093, 0.167, 1.077, -1.483]) #A_raw + A_beam
A_values_i_T = np.array([2.587, 6.147, 6.587, 2.727, -4.493, -1.193, 18.617, -3.993]) #A_raw + A_beam

dA_values_i_M = np.array([0.570, 0.570, 0.560, 0.580, 0.580, 0.580, 0.950, 0.570])
dA_values_i_T = np.array([2.870, 2.850, 2.800, 2.880, 2.900, 2.880, 4.640, 2.860])

A_values_M=np.array([A_1_M, A_2_M, A_3_M, A_4_M, A_5_M, A_6_M, A_7_M, A_8_M]) #A_raw + A_beam
A_values_T=np.array([A_1_T, A_2_T, A_3_T, A_4_T, A_5_T, A_6_T, A_7_T, A_8_T]) #A_raw + A_beam

f_pis_i_M= np.array([0.096, 0.096, 0.096, 0.096, 0.096, 0.096, 0.81, 0.096])
f_pis_i_T= np.array([0.096, 0.096, 0.096, 0.096, 0.096, 0.096, 0.81, 0.096])

f_pid_i_M = np.array([0.029, 0.029, 0.029, 0.029, 0.029, 0.029, 0.05, 0.029])
f_pid_i_T = np.array([0.029, 0.029, 0.029, 0.029, 0.029, 0.029, 0.05, 0.029])

f_NBs_i_M=np.array([0.063, 0.063, 0.063, 0.063, 0.063, 0.063, 0.51, 0.063])
f_NBs_i_T=np.array([0.063, 0.063, 0.063, 0.063, 0.063, 0.063, 0.51, 0.063])

f_NBd_i_M = np.array([0.006, 0.006, 0.006, 0.006, 0.006, 0.006, 0.09, 0.006])
f_NBd_i_T = np.array([0.006, 0.006, 0.006, 0.006, 0.006, 0.006, 0.09, 0.006])

thetas_i_M = np.array([-0.344, -0.344, -0.344, -0.344, -0.344, -0.344, -0.344, -0.344])
thetas_i_T = np.array([1.61, 1.61, 1.61, 1.61, 1.61, 1.61, 1.61, 1.61])
thetad_i_M=np.array([0.03, 0.03, 0.03, 0.03, 0.03, 0.03, 0.03, 0.03])
thetad_i_T=np.array([0.03, 0.03, 0.03, 0.03, 0.03, 0.03, 0.03, 0.03])

phis=np.array([0, pi/4, pi/2, 3*pi/4, pi, 5*pi/4, 3*pi/2, 7*pi/4])
detector=np.array([1, 2, 3, 4, 5, 6, 7, 8]) #Detector number

A_data_M=[]
for i in range(0,8):
  A_data_M.append(A_values_i_M[i] + A_BB[0] + A_bias[0])
  A_data_M[i] = np.array([A_data_M[i]])

A_data_T=[]
for i in range(0,8):
  A_data_T.append(A_values_i_T[i] + A_BB[0] + A_bias[0])
  A_data_T[i] = np.array([A_data_T[i]])

phi_M = [None]*8
phi_T = [None]*8
for i in range(0, 8):
    phi_M[i] = np.array([phis[i]]*N_M)
    phi_T[i] = np.array([phis[i]]*N_T)

dA_values_B_M = np.array([0.570*10.39, 0.570*10.39, 0.560*10.39, 0.580*10.39, 0.580*10.39, 0.580*10.39, 0.950*10.39, 0.570*10.39])
dA_values_B_T = np.array([2.870*2.07, 2.850*2.07, 2.800*2.07, 2.880*2.07, 2.900*2.07, 2.880*2.07, 4.640*2.07, 2.860*2.07])

dA_meas_f_M = [None]*8
for i in range(0, 8):
      dA_meas_f_M[i] = np.array([dA_values_B_M[i]]*(N_M))

dA_meas_f_T = [None]*8
for i in range(0, 8):
      dA_meas_f_T[i] = np.array([dA_values_B_T[i]]*(N_T))

#Adding A_BB and A_bias to A_values_f
A_meas_f_M=[]
for i in range(0,8):
  A_meas_f_M.append(A_values_M[i] + A_BB[0] + A_bias[0])

A_meas_f_T=[]
for i in range(0,8):
  A_meas_f_T.append(A_values_T[i] + A_BB[0] + A_bias[0])


# Define the number of iterations and preallocate storage for results

#num_iterations = 50  # Number of seeds/iterations
#seeds = np.arange(num_iterations)  # Array of seed numbers

#seeds = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49]
seeds = [321]
num_iterations = len(seeds)
colors = plt.cm.viridis(np.linspace(0, 1, num_iterations))

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


# Initialize lists to store results for each component across all seeds
A_fitted_B_M_results = [[] for _ in range(8)]  # 8 components
A_fitted_B_T_results = [[] for _ in range(8)]  # 8 components
A_eL_results = []
A_eT_results = []
A_piL_results = []
A_piT_results = []
dA_eL_results = []
dA_eT_results = []
dA_piL_results = []
dA_piT_results = []
all_A_eL = []
all_A_eT = []
all_A_piL = []
all_A_piT = []
mode_A_eL = []
mode_A_eT = []
mode_A_piL = []
mode_A_piT = []
correlation_matrices = []
covariance_matrices = []

#for iteration, seed in enumerate(seeds):
for iteration, (seed, color) in enumerate(zip(seeds, colors)):
  np.random.seed(seed)  # Set the random seed for reproducibility

  theta_M=np.random.normal(-0.34, 0.03)
  theta_M = np.full(N_M, theta_M)

  f_pi_M = np.random.normal(0.096, 0.029)
  f_pi_M = np.full(N_M, f_pi_M)


  f_pi7_M = np.random.normal(0.81, 0.05)
  f_pi7_M = np.full(N_M, f_pi7_M)

  f_NB_M = np.random.normal(0.063, 0.006)
  f_NB_M = np.full(N_M, f_NB_M)

  f_NB7_M = np.random.normal(0.51, 0.09)
  f_NB7_M = np.full(N_M, f_NB7_M)

  theta_T=np.random.normal(1.61, 0.03)
  theta_T = np.full(N_T, theta_T)

  f_pi_T = np.random.normal(0.096, 0.029)
  f_pi_T = np.full(N_T, f_pi_T)

  f_pi7_T = np.random.normal(0.81, 0.05)
  f_pi7_T = np.full(N_T, f_pi7_T)

  f_NB_T = np.random.normal(0.063, 0.006)
  f_NB_T = np.full(N_T, f_NB_T)

  f_NB7_T = np.random.normal(0.51, 0.09)
  f_NB7_T = np.full(N_T, f_NB7_T)

  thetas_M = np.array([theta_M, theta_M, theta_M, theta_M, theta_M, theta_M, theta_M, theta_M])
  thetas_T = np.array([theta_T, theta_T, theta_T, theta_T, theta_T, theta_T, theta_T, theta_T])
  f_pis_M = np.array([f_pi_M, f_pi_M, f_pi_M, f_pi_M, f_pi_M, f_pi_M, f_pi7_M, f_pi_M])
  f_pis_T = np.array([f_pi_T, f_pi_T, f_pi_T, f_pi_T, f_pi_T, f_pi_T, f_pi7_T, f_pi_T])
  f_NBs_M = np.array([f_NB_M, f_NB_M, f_NB_M, f_NB_M, f_NB_M, f_NB_M, f_NB7_M, f_NB_M])
  f_NBs_T = np.array([f_NB_T, f_NB_T, f_NB_T, f_NB_T, f_NB_T, f_NB_T, f_NB7_T, f_NB_T])

  asymmetry_code = """
  data {
    int<lower=0> n;
    int<lower=0> d_M;
    int<lower=0> d_T;
    matrix[n, d_M] phi_M;
    matrix[n, d_T] phi_T;
    matrix[n, d_M] theta_M;
    matrix[n, d_T] theta_T;
    matrix[n, d_M] A_meas_f_M;
    matrix[n, d_M] dA_meas_f_M;
    matrix[n, d_T] A_meas_f_T;
    matrix[n, d_T] dA_meas_f_T;
    matrix[n, d_M] f_pi_M;
    matrix[n, d_T] f_pi_T;
    matrix[n, d_M] f_NB_M;
    matrix[n, d_T] f_NB_T;
  }

  parameters {
    //component asymmetries
    real A_eL; //electron longitudinal asymmetry
    real A_eT; //electron transverse asymmetry
    real A_piL; //pion longitudinal asymmetry
    real A_piT; //pion transverse asymmetry
  }

  model {
    //priors
    A_eL ~ normal(-5.25, 100);
    A_eT ~ normal(12.3, 100);
    A_piL ~ normal(25.4, 100);
    A_piT ~ normal(-60.1, 100);
    //A_eL ~ normal(-5.25, 1.49);
    //A_eT ~ normal(12.3, 3.6);
    //A_piL ~ normal(25.4, 9.0);
    //A_piT ~ normal(-60.1, 19.3);


    //model
    for (i in 1:n)
      for (j in 1:d_M)
        A_meas_f_M[i, j] ~ normal((1 - f_NB_M[i, j]) * ((1 - f_pi_M[i, j]) * (A_eL * cos(theta_M[i, j]) + A_eT * sin(theta_M[i, j]) * sin(phi_M[i, j])) + f_pi_M[i, j] * (A_piL * cos(theta_M[i, j]) + A_piT * sin(theta_M[i, j]) * sin(phi_M[i, j]))), dA_meas_f_M[i, j]);

    for (i in 1:n)
      for (j in 1:d_T)
        A_meas_f_T[i, j] ~ normal((1 - f_NB_T[i, j]) * ((1 - f_pi_T[i, j]) * (A_eL * cos(theta_T[i, j]) + A_eT * sin(theta_T[i, j]) * sin(phi_T[i, j])) + f_pi_T[i, j] * (A_piL * cos(theta_T[i, j]) + A_piT * sin(theta_T[i, j]) * sin(phi_T[i, j]))), dA_meas_f_T[i, j]);
  }
  """

  #Put our data in a dictionary
  asym_data = {
             'n': 8,
             'd_M': N_M,
             'd_T': N_T,
             'phi_M': phi_M,
             'phi_T': phi_T,
             'theta_M': thetas_M,
             'theta_T': thetas_T,
             'A_meas_f_M': A_meas_f_M,
             'dA_meas_f_M': dA_meas_f_M,
             'A_meas_f_T': A_meas_f_T,
             'dA_meas_f_T': dA_meas_f_T,
             'f_pi_M': f_pis_M,
             'f_pi_T': f_pis_T,
             'f_NB_M': f_NBs_M,
             'f_NB_T': f_NBs_T,
            }
  #print(asym_data)
  # Compile the model
  nest_asyncio.apply()
  model = stan.build(program_code=asymmetry_code, data=asym_data, random_seed = seed)
  #model = stan.build(program_code=asymmetry_code, data=asym_data)

  # Train the model and generate samples
  asym_fit = model.sample(num_chains=1, num_samples=100000, num_warmup=500)

  df = asym_fit.to_frame()
  summary = asym_fit.to_frame().describe()

  A_eL = summary['A_eL'][1]
  A_eT = summary['A_eT'][1]
  A_piL = summary['A_piL'][1]
  A_piT = summary['A_piT'][1]

  dA_eL = summary['A_eL'][2]
  dA_eT = summary['A_eT'][2]
  dA_piL = summary['A_piL'][2]
  dA_piT = summary['A_piT'][2]

  A_eL_results.append(A_eL)
  A_eT_results.append(A_eT)
  A_piL_results.append(A_piL)
  A_piT_results.append(A_piT)

  dA_eL_results.append(dA_eL)
  dA_eT_results.append(dA_eT)
  dA_piL_results.append(dA_piL)
  dA_piT_results.append(dA_piT)

  f_NBs_M = np.array(f_NBs_M)
  f_pis_M = np.array(f_pis_M)
  thetas_M = np.array(thetas_M)
  phi_M = np.array(phi_M)

  A_fit_M = np.zeros((8, N_M))

  for i in range(8):
     for j in range(N_M):
       A_fit_M[i,j] = P * (1 - f_NBs_M[i,j]) * ((1 - f_pis_M[i,j]) * (A_eL * np.cos(thetas_M[i,j]) + A_eT * np.sin(thetas_M[i,j]) * np.sin(phi_M[i,j])) + f_pis_M[i,j] * (A_piL * np.cos(thetas_M[i,j]) + A_piT * np.sin(thetas_M[i,j]) * np.sin(phi_M[i,j])))

  f_NBs_T = np.array(f_NBs_T)
  f_pis_T = np.array(f_pis_T)
  thetas_T = np.array(thetas_T)
  phi_T = np.array(phi_T)

  A_fit_T = np.zeros((8, N_T))

  for i in range(8):
    for j in range(N_T):
       A_fit_T[i,j] = P * (1 - f_NBs_T[i,j]) * ((1 - f_pis_T[i,j]) * (A_eL * np.cos(thetas_T[i,j]) + A_eT * np.sin(thetas_T[i,j]) * np.sin(phi_T[i,j])) + f_pis_T[i,j] * (A_piL * np.cos(thetas_T[i,j]) + A_piT * np.sin(thetas_T[i,j]) * np.sin(phi_T[i,j])))

  partial_f_NB_M = -P * ((1 - f_pis_i_M) * (A_eL * np.cos(thetas_i_M) + A_eT * np.sin(thetas_i_M) * np.sin(phis)) + f_pis_i_M * (A_piL * np.cos(thetas_i_M) + A_piT * np.sin(thetas_i_M) * np.sin(phis)))
  partial_theta_rho_M = P * (1 - f_NBs_i_M) * ((1 - f_pis_i_M) * (-A_eL * np.sin(thetas_i_M) + A_eT * np.cos(thetas_i_M) * np.sin(phis)) + f_pis_i_M * (-A_piL * np.sin(thetas_i_M) + A_piT * np.cos(thetas_i_M) * np.sin(phis)))
  partial_f_pi_M = P * (1 - f_NBs_i_M) * (- (A_eL * np.cos(thetas_i_M) + A_eT * np.sin(thetas_i_M) * np.sin(phis)) + (A_piL * np.cos(thetas_i_M) + A_piT * np.sin(thetas_i_M) * np.sin(phis)))
  sigma_A_fit_M = np.sqrt((partial_f_NB_M * f_NBd_i_M)**2 + (partial_theta_rho_M * thetad_i_M)**2 + (partial_f_pi_M * f_pid_i_M)**2)

  partial_f_NB_T = -P * ((1 - f_pis_i_T) * (A_eL * np.cos(thetas_i_T) + A_eT * np.sin(thetas_i_T) * np.sin(phis)) + f_pis_i_T * (A_piL * np.cos(thetas_i_T) + A_piT * np.sin(thetas_i_T) * np.sin(phis)))
  partial_theta_rho_T = P * (1 - f_NBs_i_T) * ((1 - f_pis_i_T) * (-A_eL * np.sin(thetas_i_T) + A_eT * np.cos(thetas_i_T) * np.sin(phis)) + f_pis_i_T * (-A_piL * np.sin(thetas_i_T) + A_piT * np.cos(thetas_i_T) * np.sin(phis)))
  partial_f_pi_T = P * (1 - f_NBs_i_T) * (- (A_eL * np.cos(thetas_i_T) + A_eT * np.sin(thetas_i_T) * np.sin(phis)) + (A_piL * np.cos(thetas_i_T) + A_piT * np.sin(thetas_i_T) * np.sin(phis)))
  sigma_A_fit_T = np.sqrt((partial_f_NB_T * f_NBd_i_T)**2 + (partial_theta_rho_T * thetad_i_T)**2 + (partial_f_pi_T * f_pid_i_T)**2)

  A_fitted_B_M = np.mean(A_fit_M, axis=1)
  dA_fitted_B_M = sigma_A_fit_M

  A_fitted_B_T = np.mean(A_fit_T, axis=1)
  dA_fitted_B_T = sigma_A_fit_T

  # Store the results for each component
  for i in range(8):
     A_fitted_B_M_results[i].append(A_fitted_B_M[i])
     A_fitted_B_T_results[i].append(A_fitted_B_T[i])

  # Extract the data for the current iteration
  data_plot = {
       "A_eL": df['A_eL'],
       "A_eT": df['A_eT'],
       "A_piL": df['A_piL'],
       "A_piT": df['A_piT'],
  }
  df_iteration = pd.DataFrame(data_plot)

  # Calculate correlation and covariance matrices for the current iteration
  correlation_matrix = df_iteration.corr()
  covariance_matrix = df_iteration.cov()

  # Store matrices
  correlation_matrices.append(correlation_matrix.values)
  covariance_matrices.append(covariance_matrix.values)

  # Add values to the lists
  all_A_eL.extend(df['A_eL'])
  all_A_eT.extend(df['A_eT'])
  all_A_piL.extend(df['A_piL'])
  all_A_piT.extend(df['A_piT'])

  # Plot histograms for this iteration on the same figure
  plt.subplot(2, 2, 1)
  plt.hist(df['A_eL'], bins=20, alpha=0.5, color=color)
  plt.title("A_eL")

  plt.subplot(2, 2, 2)
  plt.hist(df['A_eT'], bins=20, alpha=0.5, color=color)
  plt.title("A_eT")

  plt.subplot(2, 2, 3)
  plt.hist(df['A_piL'], bins=20, alpha=0.5, color=color)
  plt.title("A_piL")

  plt.subplot(2, 2, 4)
  plt.hist(df['A_piT'], bins=20, alpha=0.5, color=color)
  plt.title("A_piT")


print(np.mean(A_eL_results))
print(np.mean(A_eT_results))
print(np.mean(A_piL_results))
print(np.mean(A_piT_results))

print(np.mean(dA_eL_results))
print(np.mean(dA_eT_results))
print(np.mean(dA_piL_results))
print(np.mean(dA_piT_results))

!sudo apt-get install -y fonts-liberation
!sudo fc-cache -fv
!rm ~/.cache/matplotlib -rf

# Set 'serif' as the default font family
plt.rcParams['font.family'] = 'serif'
plt.rcParams['font.serif'] = ['Times New Roman'] + plt.rcParams['font.serif']
plt.rcParams['font.size'] = 10

# Assume 'all_A_eL', 'all_A_eT', 'all_A_piL', 'all_A_piT' are lists containing all values across iterations
data = {
    "$A_e^L$": all_A_eL,
    "$A_e^T$": all_A_eT,
    "$A_{\pi}^L$": all_A_piL,
    "$A_{\pi}^T$": all_A_piT
}
df = pd.DataFrame(data)

# Setup the figure and axes for a 4x4 grid
fig, axes = plt.subplots(nrows=4, ncols=4, figsize=(15, 15))

# List of parameters
params = ['$A_e^L$', '$A_e^T$', '$A_{\pi}^L$', '$A_{\pi}^T$']

# Create a custom black colormap
black_cmap = mcolors.LinearSegmentedColormap.from_list('mycmap', ['black', 'black'])

# Plot each pair
for i, param1 in enumerate(params):
    for j, param2 in enumerate(params):
        ax = axes[i, j]
        if i == j:  # Diagonal plots
            sns.kdeplot(df[param1], ax=ax, color='black')  # Diagonal is a line KDE plot
        else:
            # Off-diagonal is a contour plot
            #sns.kdeplot(x=df[param1], y=df[param2], ax=ax, cmap=black_cmap)
            sns.kdeplot(x=df[param1], y=df[param2], ax=ax, cmap='Blues')
        ax.set_xlabel(f'Corr({param1}, {param2}) [ppm]', fontsize=14)
        ax.set_ylabel('')  # Remove y-axis label

        ax.grid(False)  # Turn off grid lines for each subplot
        if j > 0:
            ax.set_yticklabels([])  # Remove y-axis tick labels except for the first column
        if i < 3:
            ax.set_xticklabels([])  # Remove x-axis tick labels except for the last row

plt.tight_layout()
plt.show()


# Define the mapping of names to formatted LaTeX names
name_map = {
    'A_eL': r'$A_e^L$',
    'A_eT': r'$A_e^T$',
    'A_piL': r'$A_{\pi}^L$',
    'A_piT': r'$A_{\pi}^T$'
}

pairs = [
    ('A_eL', 'A_eT'),
    ('A_eL', 'A_piL'),
    ('A_eL', 'A_piT'),
    ('A_eT', 'A_piL'),
    ('A_eT', 'A_piT'),
    ('A_piL', 'A_piT')
]

# Create a 2x3 subplot
fig, axes = plt.subplots(2, 3, figsize=(10, 6))
axes = axes.flatten()  # Flatten the array of axes

# Generate each plot in the specified positions
for i, (x, y) in enumerate(pairs):
    sns.kdeplot(data=df, x=x, y=y, ax=axes[i], cmap='viridis', fill=False, levels=10)
    # Update axis labels using the name map
    axes[i].set_xlabel(name_map[x])
    axes[i].set_ylabel(name_map[y])
    axes[i].grid(False)

plt.tight_layout()
plt.show()