def check_psh_integrand_plot(k, ft_p1, muu, i1=0, i2=-1, j=[]):
"""Plot integrand of Psh to check results."""
fig, ax = utils.plot_setup(1, 1)
ax.plot(k[i1:i2], np.exp(muu * (ft_p1 - 1))[i1:i2].real)
ax.plot(k[i1:i2], np.exp(muu * (ft_p1 - 1))[i1:i2].imag)
ax.set_xscale('log')
ax.set_xlabel('k')
ax.set_ylabel(r'exp($\mu(\psi)$(P1(k)-1))')
for jj in j:
ax.axvline(fluxes[jj], color='xkcd:orchid')
return fig, ax
def check_ft_p1_plot(k, ft_p1, i1=0, i2=-1, j=[]):
"""Plot FT of P1(F) to check results."""
fig, ax = utils.plot_setup(1, 1)
ax.plot(k[i1:i2], ft_p1[i1:i2].real)
ax.plot(k[i1:i2], ft_p1[i1:i2].imag)
ax.set_xscale('log')
ax.set_xlabel('k')
ax.set_ylabel('P1(k)')
for jj in j:
ax.axvline(fluxes[jj], color='xkcd:orchid')
return fig, ax
def check_psh_plot(fluxes, psh_vals, i1=0, i2=-1, j=[]):
"""Plot Psh to check results."""
fig, ax = utils.plot_setup(1, 1)
ax.plot(fluxes[i1:i2], (fluxes * psh_vals)[i1:i2].real)
ax.plot(fluxes[i1:i2], (fluxes * psh_vals)[i1:i2].imag)
ax.set_xscale('log')
ax.set_yscale('log')
ax.set_xlabel('F')
ax.set_ylabel(r'P$_{sh}\times F$')
for jj in j:
ax.axvline(fluxes[jj], color='xkcd:orchid')
return fig, ax
def check_p1_plot(fluxes, p1_vals, i1=0, i2=-1, j=[]):
"""Plot P1 to check results."""
fig, ax = utils.plot_setup(1, 1)
ax.plot(fluxes[i1:i2], p1_vals[i1:i2])
ax.set_xscale('log')
ax.set_yscale('log')
ax.set_xlabel('F')
ax.set_ylabel('P1')
for jj in j:
ax.axvline(fluxes[jj], color='xkcd:orchid')
return fig, ax
def p1_slope_plot(params,
psi=40,
n_list=[-1, 0, 2, 4],
outfile='./output/p1_slope_plot.png'):
"""Plot slope of P1(F)."""
fig, ax = utils.plot_setup(1,
1,
figsize=(12, 8),
set_global_params=True)
n_labels = {
-1: "Som. enh.",
0: r"$s$-wave",
2: r"$p$-wave",
4: r"$d$-wave"
}
colors = iter(cm.plasma(np.linspace(0.1, 1, num=len(n_list))))
for n in n_list:
mean_params = {
'a': 77.4,
'b': 0.87 + 0.31 * n,
'c': -0.23 - 0.04 * n
}
logmin = -24
logmax = -3
fluxes = np.logspace(logmin, logmax, num=(logmax - logmin) * 20)
probs = pd.p1(fluxes, psi, mean_params=mean_params, num=200)
func = np.log10(probs)
col = next(colors)
ax.plot(fluxes[:-1], (func[:-1] - func[1:]) /
(np.log10(fluxes[:-1]) - np.log10(fluxes[1:])),
label=n_labels[n],
color=col)
ax.axhline(-1.03 / (1 + .36 * n) - 1, color=col)
print('\n')
# Plot (and save) Latent Factors
temp = xr.open_dataarray(
os.path.join(data_fp, 'sss_{}x{}.nc'.format(*dims[new_end_dim_idx])))
latent_factors = multi.model.C.detach().cpu().numpy()
if save_results:
f = open(os.path.join(save_fp, 'mrtl_latent-factors.pkl'), 'wb')
pickle.dump(latent_factors, f)
f.close()
# PLOT RESULTS
for idx in np.arange(K):
# contourf plot
max_abs = np.max(np.abs(latent_factors[:, idx])) # for scaling
fig, ax = utils.plot_setup(plot_range=[-150, -40, 10, 56])
cp = ax.contourf(temp.lon,
temp.lat,
latent_factors[:, idx].reshape(*dim),
np.linspace(-max_abs, max_abs, 15),
cmap='cmo.balance')
fig.savefig(os.path.join(save_fp, f'latent-factor{idx}_contourf.png'),
dpi=200,
bbox_inches='tight')
# plt.show()
# imshow plot
fig, ax = plt.subplots()
cp = ax.imshow(np.flipud(
latent_factors[:, idx].reshape(*dims[new_end_dim_idx])),
vmin=-max_abs,
def p1_plot(params,
psi=40,
n_list=[0, 2, 4, -1],
outfile='./output/p1_plot.png',
color=None,
shift=False,
betas=None,
fwimps=None,
residuals=True):
"""Plot P1(F)."""
if residuals is True:
fig, axs = plt.subplots(2, 1, figsize=(8, 6), sharex=True)
ax = axs[0]
else:
fig, ax = utils.plot_setup(1,
1,
figsize=(8, 6),
set_global_params=True)
n_labels = {
-1: r"Som. enh. ($n=-1$)",
0: r"$s$-wave ($n=0$)",
2: r"$p$-wave",
4: r"$d$-wave"
}
if betas is None:
betas = [params['beta']] * len(n_list)
if fwimps is None:
fwimps = params['fwimp'] * len(n_list)
if color is None:
colors = iter(cm.plasma(np.linspace(0.4, 1, num=len(n_list))))
else:
colors = iter(color)
residues = {}
for n, beta, fwimp in zip(n_list, betas, fwimps):
shiftval = 1
mean_params = {'a': 77.4, 'b': 0.87 + 0.31 * n, 'c': -0.23 - 0.04 * n}
logmin = -24
logmax = 3
fluxes = np.logspace(logmin, logmax, num=(logmax - logmin) * 20)
probs = pd.p1(fluxes,
psi,
mean_params=mean_params,
num=200,
beta=beta,
fwimp=fwimp)
# probs = [p1(flux, 40) for flux in fluxes]
func = fluxes * probs
print('integral for', n, np.trapz(func, fluxes))
if shift is True:
if n == 0:
shiftval = fluxes[func.argmax()]
print('shift is', shiftval)
else:
print('Fluxes shifted by', shiftval / fluxes[func.argmax()],
fluxes[func.argmax()])
fluxes *= shiftval / fluxes[func.argmax()]
# print(normalization)
# ax.plot(fluxes, func, label=n_labels[n] + r'$\beta$=' + str(beta), color=next(colors))
ax.plot(fluxes, func, label=n_labels[n], color=next(colors))
if residuals is True:
residues[n] = func
# print(f'slope near end for n={n}: {(func[-60]-func[-40])/(fluxes[-60]-fluxes[-40])}')
if residuals is True:
axs[-1].plot(fluxes, (residues[n_list[0]] - residues[n_list[1]]) /
residues[n_list[0]])
axs[-1].axhline(0, color='gray', lw=1)
axs[-1].set_ylabel(
rf'$(n_{n_list[0]} - n_{{{n_list[1]}}}) / n_{n_list[0]}$')
axs[-1].set_ylim(top=0.4)
ax.set_xscale('log')
ax.set_xlabel(r'Flux [photons cm$^{-2}$ yr$^{-1}$]')
ax.set_ylabel(rf'$ F \times P^n_1(F)$ at $\psi={psi}^\circ$')
ax.set_yscale('log')
# ax.set_title(rf"P_1(F) for $M_{min}={{params['M_min']:.2f}} M_\odot$, $\Psi={{psi:.2f}}^\circ$")
ax.grid()
ax.set_xticks([1e-25 * 10**i for i in range(1, 29, 2)])
ax.set_xlim(left=1e-16, right=1e-3)
ax.set_ylim(bottom=1e-8, top=1)
# lgd = ax.legend(loc='center left', bbox_to_anchor=(1.01, 0.5))
lgd = ax.legend(loc='upper right')
fig.savefig(outfile, bbox_extra_artists=(lgd, ), bbox_inches='tight')
return fig, ax
def p1_slope_plot(params,
psi=40,
n_list=[-1, 0, 2, 4],
outfile='./output/p1_slope_plot.png'):
"""Plot slope of P1(F)."""
fig, ax = utils.plot_setup(1,
1,
figsize=(12, 8),
set_global_params=True)
n_labels = {
-1: "Som. enh.",
0: r"$s$-wave",
2: r"$p$-wave",
4: r"$d$-wave"
}
colors = iter(cm.plasma(np.linspace(0.1, 1, num=len(n_list))))
for n in n_list:
mean_params = {
'a': 77.4,
'b': 0.87 + 0.31 * n,
'c': -0.23 - 0.04 * n
}
logmin = -24
logmax = -3
fluxes = np.logspace(logmin, logmax, num=(logmax - logmin) * 20)
probs = pd.p1(fluxes, psi, mean_params=mean_params, num=200)
func = np.log10(probs)
col = next(colors)
ax.plot(fluxes[:-1], (func[:-1] - func[1:]) /
(np.log10(fluxes[:-1]) - np.log10(fluxes[1:])),
label=n_labels[n],
color=col)
ax.axhline(-1.03 / (1 + .36 * n) - 1, color=col)
ax.set_xscale('log')
ax.set_xlabel(r'Flux [photons cm$^{-2}$ yr$^{-1}$]')
ax.set_ylabel('Slope of robability')
ax.set_ylim(bottom=-3, top=1)
# ax.set_xlim(left=1e-22, right=1e-2)
ax.set_title(
r"Slope of probability distribution for $M_{min}=0.01 M_\odot$, $\Psi=40^\circ$"
)
ax.grid(which='both')
ax.set_yticks(np.linspace(-3, 1, num=21))
ax.legend()
fig.savefig(outfile)
return fig, ax
