def smooth_coeff(coeff_path, cov_path, ni, nf, nmax, lmax, mmax, sn, pmass, snap, sn_out=0):
nfiles = nf-ni
S, T = coefficients_smoothing.read_coeff_matrix(coeff_path, nfiles, nmax, \
lmax, mmax, n_min=ni,\
n_max=nf, snaps=snap)
SS, TT, ST = coefficients_smoothing.read_cov_elements(cov_path, nfiles, nmax,\
lmax, mmax, n_min=ni,\
n_max=nf, snaps=snap)
if sn_out==0:
S_smooth, T_smooth, N_smooth = coefficients_smoothing.smooth_coeff_matrix(S, T, SS, TT, ST, pmass, nmax, lmax, mmax, sn, sn_out)
return S_smooth, T_smooth, N_smooth
elif sn_out==1:
S_smooth, T_smooth, N_smooth, SN_coeff = coefficients_smoothing.smooth_coeff_matrix(S, T, SS, TT, ST, pmass, nmax, lmax, mmax, sn, sn_out)
return S_smooth, T_smooth, N_smooth, SN_coeff
def load_scf_coefficients(coeff_files, nmax, lmax, mmax, pmass, sn):
"""
Load coefficients.
## this is the s/n sample to take the mean values of the coefficients!
min_sample: n_min = 0
max_sample: nmax = 1
# pmass : particle mass
# sn
"""
#ncoeff_sample = max_sample - min_sample
# smoothing routines -------------------
data = np.loadtxt(coeff_files)
S = data[:, 0]
SS = data[:, 1]
T = data[:, 2]
TT = data[:, 3]
ST = data[:, 4]
S = coefficients_smoothing.reshape_matrix(S, nmax, lmax, lmax)
SS = coefficients_smoothing.reshape_matrix(SS, nmax, lmax, lmax)
T = coefficients_smoothing.reshape_matrix(T, nmax, lmax, lmax)
TT = coefficients_smoothing.reshape_matrix(TT, nmax, lmax, lmax)
ST = coefficients_smoothing.reshape_matrix(ST, nmax, lmax, lmax)
S_smooth, T_smooth, N_smooth = coefficients_smoothing.smooth_coeff_matrix(S, T, SS, TT, ST,\
pmass,\
nmax,\
lmax,\
mmax,\
sn,\
sn_out=0)
return S_smooth, T_smooth
def load_scf_coefficients(coeff_files, cov_files, nmax, lmax, mmax, min_sample,
max_sample, pmass, sn):
"""
Load coefficients.
TODO : REMOVE THIS FUNCTION WHEN SCRIPT IS FULLY WORKING
"""
ncoeff_sample = max_sample - min_sample
print(pmass)
S, T = coefficients_smoothing.read_coeff_matrix(coeff_files,
ncoeff_sample,
nmax,
lmax,
mmax,
min_sample,
max_sample,
snaps=90)
SS, TT, ST = coefficients_smoothing.read_cov_elements(cov_files,
ncoeff_sample,
nmax,
lmax,
mmax,
min_sample,
max_sample,
snaps=90)
S_smooth, T_smooth, N_smooth = coefficients_smoothing.smooth_coeff_matrix(
S, T, SS, TT, ST, pmass, nmax, lmax, mmax, sn, sn_out=0)
return S_smooth, T_smooth
def acceleration_plot(S, T, SS, TT, ST, figname, cbar_name):
sn = [0, 1, 2, 3, 4, 5]
fig = plt.figure(figsize=(10, 14))
for i in range(len(sn)):
S_smooth, T_smooth, N_smooth = coefficients_smoothing.smooth_coeff_matrix(
S, T, SS, TT, ST, mass, 20, 20, 20, sn[i])
a_biff = biff.gradient(np.ascontiguousarray(xyz),
S_smooth,
T_smooth,
M=1,
r_s=40.85,
G=1)
a_biff_all = np.sqrt(a_biff[:, 0]**2 + a_biff[:, 1]**2 +
a_biff[:, 2]**2)
a_biff_0 = biff.gradient(np.ascontiguousarray(xyz),
np.array([[[S_smooth[0, 0, 0]], [0], [0]]]).T,
np.array([[[T_smooth[0, 0, 0]], [0], [0]]]).T,
M=1,
r_s=40.85,
G=1)
a_biff_all_0 = np.sqrt(a_biff_0[:, 0]**2 + a_biff_0[:, 1]**2 +
a_biff_0[:, 2]**2)
plt.subplot(3, 2, i + 1)
#levels = np.arange(-0.03, 0.12, 0.005)
im = plt.contourf(y_grid,
z_grid, ((a_biff_all / a_biff_all_0) - 1).reshape(
bins, bins),
40,
origin='lower',
cmap='viridis')
plt.text(-180, 160, 'Ncoeff={}'.format(N_smooth), color='w')
plt.text(-180, 130, 'S/N={}'.format(sn[i]), color='w')
plt.xlim(-200, 200)
plt.ylim(-200, 200)
if ((i == 4) | (i == 5)):
plt.xlabel('y[kpc]')
if (i + 1) % 2 == 1:
plt.ylabel('z[kpc]')
cb_ax = fig.add_axes([0.93, 0.1, 0.02, 0.8])
cbar = plt.colorbar(im, cax=cb_ax)
fig.suptitle(
'Relative acceleration of MW + LMC unbound particles {}'.format(
cbar_name),
y=0.93)
cbar.set_label('$\Delta |a|$')
plt.savefig(figname + '.pdf', bbox_inches='tight')
plt.savefig(figname + '.png', bbox_inches='tight')
return 0
def load_scf_coefficients_cov(coeff_files, nmax, lmax, mmax, min_sample,
max_sample, pmass, sn):
"""
Load smoothed coefficients
Parameters:
=========
coeff_files : str
path to coefficients files
nmax : int
nmax order of the expansion
lmax : int
lmax order of the expansion
mmax : int
mmax order of the expansion
min_sample : int
minimum value of variance sample
max_sample : int
maximum value of variance sample
pmass : float
particle mass
sn : float
signal-to-noise
Return:
=======
Ssmooth : smoothed coefficients
Tsmooth : smoothed coefficients
"""
ncoeff_sample = max_sample - min_sample
print(pmass)
S, T, SS, TT, ST = coefficients_smoothing.read_coeffcov_matrix(
coeff_files,
ncoeff_sample,
nmax,
lmax,
mmax,
min_sample,
max_sample,
snaps=0)
S_smooth, T_smooth, N_smooth = coefficients_smoothing.smooth_coeff_matrix(
S, T, SS, TT, ST, pmass, nmax, lmax, mmax, sn)
return S_smooth, T_smooth
def relative_rho(S, T, SS, TT, ST, figname, title_name):
sn = [0, 1, 2, 3, 4, 5]
fig = figure(figsize=(10, 14))
for i in range(len(sn)):
S_smooth, T_smooth, N_smooth = coefficients_smoothing.smooth_coeff_matrix(
S, T, SS, TT, ST, mass, 20, 20, 20, sn[i])
rho_biff = biff.density(np.ascontiguousarray(xyz),
S_smooth,
T_smooth,
M=1,
r_s=40.85)
print(S_smooth[0, 0, 0])
rho_biff_0 = biff.density(np.ascontiguousarray(xyz),
np.array([[[S_smooth[0, 0, 0]], [0],
[0]]]).T,
np.array([[[T_smooth[0, 0, 0]], [0],
[0]]]).T,
M=1,
r_s=40.85)
subplot(3, 2, i + 1)
#levels = np.arange(-14, -2.4, 0.5)
im = contourf(y_grid,
z_grid,
((rho_biff / rho_biff_0) - 1).reshape(bins, bins),
30,
origin='lower',
cmap='Spectral_r')
text(-180, 160, 'Ncoeff={}'.format(N_smooth), color='k')
text(-180, 130, 'S/N={}'.format(sn[i]), color='k')
xlim(-200, 200)
ylim(-200, 200)
if ((i == 4) | (i == 5)):
xlabel('y[kpc]')
if (i + 1) % 2 == 1:
ylabel('z[kpc]')
cb_ax = fig.add_axes([0.93, 0.1, 0.02, 0.8])
cbar = colorbar(im, cax=cb_ax)
fig.suptitle(
'Relative density of MW + LMC unbound particles {}'.format(title_name),
y=0.93)
cbar.set_label(r'$\Delta \rho$')
savefig(figname + '.pdf', bbox_inches='tight')
savefig(figname + '.png', bbox_inches='tight')
def rho_square(i, sn):
"""
Computes the `True' density square as defined
in equation 5 of the companion doc.
It use the coefficients computed with the particle distribution
in each particle batch i.
"""
path = '/home/u9/jngaravitoc/codes/Kullback-Leibler/data/'
t1 = time_now()
S, T = coefficients_smoothing.read_coeff_matrix(
path + 'mwlmc_hal_sn_test_coeff_sample_',
1,
nmax,
lmax,
lmax,
i,
i + 1,
snaps=1)
SS, TT, ST = coefficients_smoothing.read_cov_elements(
path + 'mwlmc_hal_sn_test_covmat_sample_',
1,
nmax,
lmax,
lmax,
i,
i + 1,
snaps=1)
S_smooth, T_smooth, N_coeff = coefficients_smoothing.smooth_coeff_matrix(
S, T, SS, TT, ST, mass, nmax, lmax, lmax, sn_range[sn])
rho2 = coefficients_sum_fast(S_smooth, T_smooth, rs)
t2 = time_now()
time_diff(t1, t2)
return rho2
def rho_plots(S, T, SS, TT, ST, figname, cbar_name):
sn = [0, 1, 2, 3, 4, 5]
fig = plt.figure(figsize=(10, 14))
for i in range(len(sn)):
S_smooth, T_smooth, N_smooth = coefficients_smoothing.smooth_coeff_matrix(
S, T, SS, TT, ST, mass, 20, 20, 20, sn[i])
rho_biff = biff.density(np.ascontiguousarray(xyz),
S_smooth,
T_smooth,
M=1,
r_s=40.85)
plt.subplot(3, 2, i + 1)
#levels = np.arange(-14, -2.4, 0.5)
im = plt.contourf(
y_grid,
z_grid,
np.log10(np.abs(rho_biff).reshape(bins, bins)),
40,
origin='lower',
cmap='Spectral_r',
)
plt.text(-180, 160, 'Ncoeff={}'.format(N_smooth), color='k')
plt.text(-180, 130, 'S/N={}'.format(sn[i]), color='k')
plt.xlim(-200, 200)
plt.ylim(-200, 200)
if ((i == 4) | (i == 5)):
plt.xlabel('y[kpc]')
if (i + 1) % 2 == 1:
plt.ylabel('z[kpc]')
cb_ax = fig.add_axes([0.93, 0.1, 0.02, 0.8])
cbar = plt.colorbar(im, cax=cb_ax)
fig.suptitle('Density of MW + LMC unbound particles {}'.format(cbar_name),
y=0.93)
cbar.set_label(r'$\rm{Log_{10}} \rho$')
plt.savefig(figname + '.pdf', bbox_inches='tight')
plt.savefig(figname + '.png', bbox_inches='tight')
return 0
def kpq_sn_estimator(i):
pos = np.loadtxt('./MW_100M_b1_dm_part_1e6_300.txt')
N_particles = len(pos)
n_batches = 100
particles_in_batch = int(N_particles / n_batches)
sn_range = np.arange(0, 11, 0.2)
Hp = np.zeros(len(sn_range))
Kpq = np.zeros(len(sn_range))
Kpq_rho = np.zeros(len(sn_range))
N_part = 990000 # len(pos_grid)
m_p = 1 / N_part
rho_factor = 1e4 / N_part
m_p_sim = 1E-4
N_coeff = np.zeros(len(sn_range))
path = '/home/xzk/work/github/MW-LMC-SCF/code/data_KL_SN/'
S, T = coefficients_smoothing.read_coeff_matrix(
path + 'mwlmc_hal_sn_test_coeff_sample_',
i + 1,
20,
20,
20,
0,
i + 1,
snaps=1)
SS, TT, ST = coefficients_smoothing.read_cov_elements(
path + 'mwlmc_hal_sn_test_covmat_sample_',
i + 1,
20,
20,
20,
0,
i + 1,
snaps=1)
S_0, T_0, N_smooth_0 = coefficients_smoothing.smooth_coeff_matrix(
S, T, SS, TT, ST, m_p_sim, 20, 20, 20, 0) # sn
xyz1 = np.array([
pos[:i * (particles_in_batch), 0], pos[:i * (particles_in_batch), 1],
pos[:i * (particles_in_batch), 2]
]).T
xyz2 = np.array([
pos[(i + 1) * (particles_in_batch):,
0], pos[(i + 1) * (particles_in_batch):, 1],
pos[(i + 1) * (particles_in_batch):, 2]
]).T
xyz = np.concatenate((xyz1, xyz2))
assert (len(xyz == 990000))
rho_reference = biff.density(np.ascontiguousarray(xyz.astype(float)),
S_0,
T_0,
M=1,
r_s=40.85)
print('Here')
"""
for sn in range(len(sn_range)):
S_smooth, T_smooth, N_coeff[sn] = coefficients_smoothing.smooth_coeff_matrix(S, T, SS, TT, ST, m_p_sim, 20, 20, 20, sn_range[sn])
rho_estimate = biff.density(np.ascontiguousarray(xyz.astype(float)), S_smooth, T_smooth, M=1, r_s=40.85)
Hp[sn] = np.sum(m_p*np.log(rho_factor*np.abs(rho_estimate)))
Kpq[sn] = np.sum(m_p*np.log(rho_factor*np.abs(rho_reference))) - Hp[sn]
Kpq_rho[sn] = np.sum(rho_factor*rho_reference*np.log(np.abs(rho_reference/rho_estimate)))
"""
return Hp, Kpq, Kpq_rho, N_coeff
