def plot_2d_pdf(ax, pdf, mass=None, age=None):
ax.minorticks_on()
resolution = len(pdf)
cmap = palettable.cmocean.sequential.Ice_20.mpl_colormap
cmap = palettable.scientific.sequential.Oslo_3.mpl_colormap
im = ax.imshow(pdf.T + 1,
norm=LogNorm(),
alpha=.75,
interpolation="bilinear",
extent=extent,
origin="lower",
cmap=cmap)
if mass is not None and age is not None:
ax.scatter(mass,
age,
s=150,
marker='*',
c="#ff4747",
alpha=1.0,
zorder=10)
ax.set_xlabel(r"$m_{\textsc{wdm}}$")
ax.set_xticks([10, 20, 30, 40, 50])
ax.set_ylabel("Stream age in Gyr")
ax.grid(True, which="both", alpha=.15, zorder=0, color="white")
make_square(ax)
def display(ax_trace, ax_density, theta_index=1):
# Trace
ax_trace.minorticks_on()
ax_trace.plot(range(num_samples),
chain.samples.numpy(),
color="black",
lw=2)
ax_trace.set_xlim([0, num_samples])
ax_trace.set_xticks([])
ax_trace.set_ylabel(r"$\theta_" + str(theta_index) + "$")
limits = ax_trace.get_ylim()
# Density
ax_density.minorticks_on()
ax_density.hist(chain.samples.numpy(),
bins=50,
lw=2,
color="black",
histtype="step",
density=True)
ax_density.yaxis.tick_right()
ax_density.yaxis.set_label_position("right")
ax_density.set_ylabel("Probability mass function")
ax_density.set_xlabel(r"$\theta_" + str(theta_index) + "$")
ax_density.set_xlim(limits)
# Aspects
make_square(ax_density)
ax_trace.set_aspect("auto")
ax_trace.set_position([0, 0, .7, 1])
ax_density.set_position([.28, 0, 1, 1])
def plot_1d_pdf(ax, pdf, resolution=default_resolution):
masses = np.linspace(prior_1d.low.item(), prior_1d.high.item(), resolution)
ax.minorticks_on()
ax.plot(masses, pdf, lw=2, color="black")
ax.set_xlabel(r"$m_\textsc{wdm}$")
ax.set_xticks(masses_xticks)
ax.grid(True, which="both", alpha=.15)
make_square(ax)
def plot_1d_profile_likelihood(ax, profile_likelihood):
resolution = len(profile_likelihood)
masses = np.linspace(prior_1d.low.item(), prior_1d.high.item(), resolution)
ax.minorticks_on()
ax.plot(masses, profile_likelihood, lw=2, color="black")
ax.set_xlabel(r"$m_\textsc{wdm}$")
ax.set_xticks(masses_xticks)
ax.set_ylabel(r"Profile likelihood")
ax.set_ylabel(
r"$-2\mathbb{E}\left[\log r(x\vert\vartheta) - \log r(x\vert\hat{\vartheta})\right]$"
)
ax.grid(True, which="both", alpha=.15)
make_square(ax)
def plot_losses(losses_train,
losses_test,
epochs=None,
log=True,
figsize=None,
ylim=None):
with torch.no_grad():
# Check if the losses have to be converted to log-space.
if log:
ylabel = "Logarithmic loss"
losses_train = losses_train.log()
losses_test = losses_test.log()
else:
ylabel = "Loss"
# Check if custom epochs have been specified.
if epochs is None:
if losses_train.dim() == 2:
num_epochs = losses_train.shape[1]
else:
num_epochs = losses_train.shape[0]
epochs = np.arange(1, num_epochs + 1)
# Allocate a figure with shared y-axes.
figure, axes = plt.subplots(nrows=1,
ncols=2,
figsize=figsize,
sharey=True)
plot_loss(axes[0],
losses_train,
epochs=epochs,
title="Training loss",
xlabel="Epochs",
ylabel=ylabel)
plot_loss(axes[1],
losses_test,
epochs=epochs,
title="Test loss",
xlabel="Epochs",
ylabel=None)
figure.tight_layout()
if ylim is not None:
axes[0].set_ylim(ylim)
axes[1].set_ylim(ylim)
make_square(axes[0])
make_square(axes[1])
return figure
def plot(paths, title=None):
# Prepare the data
data = stack(paths)
figure, ax = plt.subplots(1)
mean = data.mean(dim=0)
std = data.std(dim=0)
# Plot the data
epochs = np.arange(1, len(mean) + 1)
ax.set_title(title)
ax.plot(epochs, mean, lw=2, color="black")
ax.fill_between(epochs, mean - std, mean + std, color="black", alpha=0.1)
ax.minorticks_on()
ax.set_xlabel("Epochs")
ax.set_ylabel("Loss")
make_square(ax)
return figure
def plot_loss(ax, losses, epochs=None, title=None, xlabel=None, ylabel=None):
with torch.no_grad():
# Check if the standard deviation needs to be computed.
if losses.dim() == 2 and losses.shape[1] > 1:
mean, std = losses.mean(dim=0), losses.std(dim=0)
ax.plot(epochs, mean, color="black", lw=2, label="Loss")
ax.fill_between(epochs,
mean - std,
mean + std,
alpha=.25,
color="black",
label=r"$\pm1\sigma$ loss")
ax.legend()
else:
ax.plot(epochs, losses.numpy(), color="black", lw=2)
ax.set_title(title)
ax.set_xlabel(xlabel)
ax.set_ylabel(ylabel)
ax.minorticks_on()
make_square(ax)
def plot_autocorrelation(chain, interval=2, max_lag=100, radius=1.1):
if max_lag is None:
max_lag = chain.size()
autocorrelations = chain.autocorrelations()[:max_lag]
lags = np.arange(0, max_lag, interval)
autocorrelations = autocorrelations[lags]
plt.ylim([-radius, radius])
center = .5
for index, lag in enumerate(lags):
autocorrelation = autocorrelations[index]
plt.axvline(lag,
center,
center + autocorrelation / 2 / radius,
c="black")
plt.xlabel("Lag")
plt.ylabel("Autocorrelation")
plt.minorticks_on()
plt.axhline(0, linestyle="--", c="black", alpha=.75, lw=2)
make_square(plt.gca())
figure = plt.gcf()
return figure
losses = []
for epoch in range(epochs):
data_loader = torch.utils.data.DataLoader(dataset,
batch_size=batch_size,
drop_last=True)
num_batches = len(data_loader)
data_loader = iter(data_loader)
for batch_index in range(num_batches):
optimizer.zero_grad()
inputs, outputs = next(data_loader)
inputs = inputs.to(hypothesis.accelerator)
outputs = outputs.to(hypothesis.accelerator)
loss = criterion(inputs=inputs, outputs=outputs)
loss.backward()
optimizer.step()
losses.append(loss.item())
losses = np.array(losses)
plt.plot(np.log(losses), lw=2, color="black")
plt.minorticks_on()
plt.xlabel("Gradient updates")
plt.ylabel("Logarithmic loss")
make_square(plt.gca())
plt.show()
ratio_estimator = ratio_estimator.cpu()
ratio_estimator.eval()
## POSTERIOR INFERENCE
def plot_1d_pdf_std(ax, pdf, std, resolution=default_resolution):
masses = np.linspace(prior_1d.low.item(), prior_1d.high.item(), resolution)
ax.fill_between(masses, pdf - std, pdf + std, color="black", alpha=.15)
make_square(ax)
def plot_stream(ax, phi, stream):
ax.step(phi, stream.reshape(-1), lw=2, color="black")
ax.set_ylim([0, 2])
ax.minorticks_on()
make_square(ax)