def scf_density_grid_fast(x,
y,
z,
Smw,
Tmw,
Slmc,
Tlmc,
nbins,
rs_mw,
rs_lmc,
lmc_com,
quantity,
G=1):
q_all = np.zeros((nbins, nbins, nbins))
xyz_lmc = np.array([x - lmc_com[0], y - lmc_com[1], z - lmc_com[2]]).T
xyz = np.array([x, y, z]).T
print(len(xyz_lmc[:, 0]))
print(xyz_lmc[:, 0])
if quantity == "density":
q_all = biff.density(np.ascontiguousarray(xyz_lmc), Slmc, Tlmc, M=1, r_s=rs_lmc) \
+ biff.density(np.ascontiguousarray(xyz), Smw, Tmw, M=1, r_s=rs_mw)
if quantity == "potential":
q_all = biff.potential(np.ascontiguousarray(xyz_lmc), Slmc, Tlmc, M=1, r_s=rs_lmc, G=G) \
+ biff.potential(np.ascontiguousarray(xyz), Smw, Tmw, M=1, r_s=rs_mw, G=G)
if quantity == "acceleration":
almc = biff.gradient(xyz_lmc, Slmc, Tlmc, M=1, r_s=rs_lmc, G=G)
amw = biff.gradient(xyz, Smw, Tmw, M=1, r_s=rs_mw, G=G)
q_all = np.sqrt(np.sum(almc**2, axis=1)) + np.sqrt(
np.sum(amw**2, axis=1))
return q_all
def combine_bfe_rho(S1, T1, S2, T2, y_grid, z_grid, lmc_com, nbins):
rho_mwlmc = np.zeros((nbins, nbins))
for i in range(nbins):
for j in range(nbins):
rho_mwlmc[i][j] = biff.density(np.array([[0-lmc_com[0]],
[y_grid[0][i]-lmc_com[1]],
[z_grid[j,0]-lmc_com[2]]]).T,
S2, T2, M=1,r_s=10) + \
biff.density(np.array([[0], [y_grid[0][i]], [z_grid[j,0]]]).T,
S1, T1, M=1, r_s=40.85)
return rho_mwlmc
def _ln_potential_prior(self, pars):
lp = 0.
lp += -(pars['Snlm']**2).sum()
grad = biff.gradient(self._xyz,
Snlm=pars['Snlm'],
Tnlm=pars['Tnlm'],
nmax=self.nmax,
lmax=0,
G=_G,
M=pars['m'],
r_s=pars['r_s'])
if np.any(grad < 0.):
return -np.inf
dens = biff.density(self._xyz,
Snlm=pars['Snlm'],
Tnlm=pars['Tnlm'],
nmax=self.nmax,
lmax=0,
M=pars['m'],
r_s=pars['r_s'])
if np.any(dens < 0.):
return -np.inf
return lp
def scf_density_grid(x,
y,
z,
Smw,
Tmw,
Slmc,
Tlmc,
nbins,
rs_mw,
rs_lmc,
lmc_com,
quantity,
G=1):
q_all = np.zeros((nbins, nbins, nbins))
for i in range(nbins):
for j in range(nbins):
for k in range(nbins):
# These line is flipping x with y!
xyz_lmc = np.array([[x[0, i, 0] - lmc_com[0]],
[y[j, 0, 0] - lmc_com[1]],
[z[0, 0, k] - lmc_com[2]]]).T
xyz = np.array([[x[0, i, 0]], [y[j, 0, 0]], [z[0, 0, k]]]).T
if quantity == "density":
q_all[i][j][k] = \
biff.density(xyz_lmc, Slmc, Tlmc, M=1, r_s=rs_lmc) \
+ biff.density(xyz, Smw, Tmw, M=1, r_s=rs_mw)
if quantity == "potential":
q_all[i][j][k] = \
biff.potential(np.ascontiguousarray(xyz_lmc), Slmc, Tlmc, M=1, r_s=rs_lmc,
G=G) \
+ biff.potential(np.ascontiguousarray(xyz), Smw, Tmw, M=1, r_s=rs_mw, G=G)
if quantity == "acceleration":
almc = biff.gradient(xyz_lmc,
Slmc,
Tlmc,
M=1,
r_s=rs_lmc,
G=G)
amw = biff.gradient(xyz, Smw, Tmw, M=1, r_s=rs_mw, G=G)
q_all[i][j][k] = np.sqrt(np.sum(almc**2)) \
+ np.sqrt(np.sum(amw**2))
return q_all.flatten()
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_bfe_all(self, S, T, S_mw, T_mw, S_mw_wake, T_mw_wake, S_lmc, T_lmc,
S_mwlmc, T_mwlmc):
xyz, xyz_shift = self.grid_cartessian()
# Monopole
rho_ref = biff.density(xyz, S_mw_ref, T_mw_ref, M=1, r_s=self.rs)
rho_wake = biff.density(xyz, S_mw, T_mw, M=1, r_s=self.rs)
rho_mwwake = biff.density(xyz, S_mw_wake, T_mw_wake, M=1, r_s=self.rs)
rho_lmcshift = biff.density(xyz_shift,
S_lmc,
T_lmc,
M=1,
r_s=self.rs_sat)
#rho_lmcshift = biff.density(xyz, S_lmc, T_lmc, M=1, r_s=self.rs_sat)
# MW + Debris
rho_mwlmc = biff.density(xyz, S_mwlmc, T_mwlmc, M=1, r_s=self.rs)
return rho_ref, rho_mwwake, rho_wake, rho_lmcshift, rho_mwlmc
def scf_density(x_grid, y_grid, z_grid, S, T, r_s_mw):
xyz = np.ascontiguousarray(
np.double(
np.array([x_grid.flatten(),
y_grid.flatten(),
z_grid.flatten()]).T))
dens_ratio_all = biff.density(xyz, S, T, M=1, r_s=r_s_mw)
return dens_ratio_all
def _ln_potential_prior(self, pars):
lp = 0.
lp += -(pars['Snlm']**2).sum()
grad = biff.gradient(self._xyz, Snlm=pars['Snlm'], Tnlm=pars['Tnlm'],
nmax=self.nmax, lmax=0, G=_G, M=pars['m'], r_s=pars['r_s'])
if np.any(grad < 0.):
return -np.inf
dens = biff.density(self._xyz, Snlm=pars['Snlm'], Tnlm=pars['Tnlm'],
nmax=self.nmax, lmax=0, M=pars['m'], r_s=pars['r_s'])
if np.any(dens < 0.):
return -np.inf
return lp
def BFE_density(Snlm, Tnlm, true_M, true_r_s, xmin, xmax, ymin, ymax, nbinsx,
nbinsy, z_plane, **kwargs):
"""
Computes the density profile of the BFE
"""
y_bins = np.linspace(xmin, xmax, nbinsx)
z_bins = np.linspace(ymin, ymax, nbinsy)
y, z = np.meshgrid(y_bins, z_bins)
pos = np.array(
[np.ones(len(y.flatten())) * z_plane,
y.flatten(),
z.flatten()]).T
rho_bfe = biff.density(np.ascontiguousarray(pos.astype(np.double)), Snlm,
Tnlm, true_M, true_r_s)
if 'S2' in kwargs.keys():
S2 = kwargs['S2']
T2 = kwargs['T2']
M2 = kwargs['M2']
r_s2 = kwargs['r_s2']
ylmc = kwargs['ylmc']
zlmc = kwargs['zlmc']
print(ylmc, zlmc)
y_bins2 = np.linspace(ylmc - xmin, ylmc - xmax, nbinsx)
z_bins2 = np.linspace(zlmc - ymin, zlmc - ymax, nbinsy)
y2, z2 = np.meshgrid(y_bins2, z_bins2)
pos2 = np.array(
[np.ones(len(y2.flatten())) * z_plane,
y2.flatten(),
z2.flatten()]).T
rho_bfe2 = biff.density(np.ascontiguousarray(pos2.astype(np.double)),
S2, T2, M2, r_s2)
return rho_bfe, rho_bfe2
def BFE_density_profile(pos, Snlm, Tnlm, true_M, true_r_s, rmax, nbins=20):
"""
Computes the density profile of the BFE
"""
#r = np.logspace(-2, rmax, 512)
#pos = np.zeros((len(r), 3))
#pos[:,0] = r
rho_bfe = biff.density(np.ascontiguousarray(pos.astype(np.double)), Snlm,
Tnlm, true_M, true_r_s)
r = (pos[:, 0]**2 + pos[:, 1]**2 + pos[:, 2]**2)**0.5
r_bins = np.linspace(0.01, rmax, nbins)
rho_bins = np.zeros(len(r_bins))
for i in range(len(rho_bins) - 1):
index_c = np.where((r < r_bins[i + 1]) & (r >= r_bins[i]))[0]
rho_bins[i] = np.mean(rho_bfe[index_c])
dr = (r_bins[1] - r_bins[0]) / 2.
return r_bins + dr, rho_bins
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
