# %%
log_abs_det_jacobian(flow[-1], torch.Tensor([[1, 1]]), None)
# %%
flow[-1].log_abs_det_jacobian(torch.Tensor([[1, 1]]), None)
# %%
x = torch.Tensor(X0)
# %%
for i, bij in enumerate(flow):
if i % 2 == 0:
print(bij.log_abs_det_jacobian(torch.Tensor([[1, 1]]), None))
# %%
-flow.log_prob(torch.Tensor([[1, 1]]))
# %%
_, log_abs_det = flow.forward(torch.Tensor([[1, 1]]))
log_abs_det
# %%
flow.log_prob(x_samples).mean()
# %%
z_samples = flow.base_dist.sample((1000,))
x_samples, log_abs_det_fw = flow.forward(z_samples)
log_abs_det_fw
# %%
best_loss = torch.Tensor([float("+inf")])
attempts = 0
flow.train()
for epoch in trange(n_epochs):
batches = range((len(X) - 1) // bs + 1)
for i in batches:
start_i = i * bs
end_i = start_i + bs
xb = X[start_i:end_i]
it = epoch * len(batches) + i + 1
opt.zero_grad()
loss = -flow.log_prob(xb).mean()
#if loss <= 0:
# if attempts < 100:
# attempts += 1
# continue
# else:
# print("Loss has diverged, halting train and not backpropagating")
# break
if loss <= best_loss:
best_loss = loss
best_params = flow.state_dict()
loss.backward()
writer = SummaryWriter(f"./tensorboard_logs/{now_str()}")
best_loss = torch.Tensor([float("+inf")])
#attempts = 0
for epoch in trange(n_epochs):
batches = range((len(X) - 1) // bs + 1)
for i in batches:
start_i = i * bs
end_i = start_i + bs
xb = X[start_i:end_i]
it = epoch * len(batches) + i + 1
opt.zero_grad()
loss = -flow.log_prob(xb).mean()
#if loss <= 0:
# if attempts < 100:
# attempts += 1
# continue
# else:
# print("Loss has diverged, halting train and not backpropagating")
# break
if loss <= best_loss:
best_loss = loss
best_flow_params = flow.state_dict()
loss.backward()
opt.step()
cb = ax.contourf(xx, yy, zz, 50, cmap="rainbow")
ax.set_aspect("equal")
divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="5%", pad=0.05)
plt.colorbar(cb, cax=cax)
plt.show()
# %%
x = np.linspace(-4, 4, 1000)
z = np.array(np.meshgrid(x, x)).transpose(1, 2, 0)
z = np.reshape(z, [z.shape[0] * z.shape[1], -1])
with torch.no_grad():
densities = flow.log_prob(torch.Tensor(z)).exp().numpy()
#densities = flow.log_prob(torch.Tensor(z)).numpy()
mesh = z.reshape([1000, 1000, 2]).transpose(2, 0, 1)
xx = mesh[0]
yy = mesh[1]
f, ax = plt.subplots(figsize=(10, 10))
zz = densities.reshape([1000, 1000])
cb = ax.contourf(xx, yy, zz, 50, cmap="rainbow")
ax.set_aspect("equal")
divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="5%", pad=0.05)