<jupyter_start><jupyter_text>Train flow to transform gaussian distr into moon distr<jupyter_code>state_dict_path = Path('flow')
if state_dict_path.exists():
    flow.load_state_dict(torch.load(state_dict_path))
    moon_data, _ = make_moons(n_samples=BATCHSIZE, noise=0.05)
else:
    optimizer = torch.optim.Adam(flow.parameters(), lr=0.001)

    for i in range(1000):
        optimizer.zero_grad()
        # sample data from the moons distribution
        moon_data, label = make_moons(n_samples=BATCHSIZE, noise=0.05)
        x = torch.Tensor(moon_data)
        # pass to INN and get transformed variable z and log Jacobian determinant
        z, log_jac_det = flow(x)
        # calculate the negative log-likelihood of the model with a standard normal prior
        loss = 0.5 * torch.sum(z ** 2, 1) - log_jac_det
        loss = loss.mean() / N_DIM
        # backpropagate and update the weights
        loss.backward()
        optimizer.step()
    torch.save(flow.state_dict(), 'flow')
plt.title('Moon Data')
plt.scatter(moon_data[:, 0], moon_data[:, 1])
plt.show()
q1 = torch.randn(BATCHSIZE, N_DIM)
q2 = torch.randn(BATCHSIZE, N_DIM)
plt.title('Gauss Data = $q_1$')
plt.scatter(q1[:, 0], q1[:, 1])
plt.show()
z1, _ = flow(q1, rev=True)
z2, _ = flow(q2, rev=True)
z1_detached = z1.detach().numpy()
plt.title('$z_1 = f(q_1)$')
plt.scatter(z1_detached[:, 0], z1_detached[:, 1])
plt.show()


# z_tilde = torch.mean(torch.cat((z1.unsqueeze(dim=0), z2.unsqueeze(dim=0))), dim=0)
z_tilde_detached = (z1 + z2).detach().numpy()
plt.title(r'$\tilde{z} = f(q_1) + f(q_2)$')
plt.scatter(z_tilde_detached[:, 0], z_tilde_detached[:, 1])
plt.show()
q12 = q1 + q2
plt.title('$q_1$ + $q_2$')
plt.scatter(q12[:, 0], q12[:, 1])
plt.show()
z, _ = flow(q12, rev = True)
z = z.detach().numpy()
plt.title(r'$z = f^{-1}(\tilde{z}) = f^{-1}(MoE(f(q_1), f(q_2)))$')
plt.scatter(z[:, 0], z[:, 1])
plt.show()
q1, _ = flow(z1)
q1 = q1.detach().numpy()
plt.title(r'$q_1 = f^{-1}(z1)$')
plt.scatter(q1[:, 0], q1[:, 1])
plt.show()<jupyter_output><empty_output>
StarcoderdataJupyterStructuredCleanDedup