Planar flow transformed distribution does not integrate to 1

Misc.
1

Just trying to generate the pdf of a unit Gaussian transformed by the Planar Flow with the code below.

But seems like the transformed pdf does not integrate to 1.

Perhaps I am doing something wrong with the way I am computing the transformed pdf.

import numpy as np
import torch as torch
import torch.nn as nn
import matplotlib.pylab as plt
import pyro.distributions as dist
import pyro.distributions.transforms as trans

# define a meshgrid

x1 = torch.tensor(data=np.linspace(-5,5,100))
x2 = torch.tensor(data=np.linspace(-5,5,100))

x1_s, x2_s = torch.meshgrid(x1, x2)
x_field = torch.tensor(np.concatenate([x1_s[..., None], x2_s[..., None]], axis=-1)).float()

# base distribution 
base_dist = dist.MultivariateNormal(loc=torch.zeros(2), covariance_matrix=torch.eye(2))

# plot base

plt.figure(figsize=(5,5))
plt.contourf(x1_s, x2_s, torch.exp(base_dist.log_prob(x_field)))

# Check if integrates to 1

print(np.trapz(np.trapz(torch.exp(base_dist.log_prob(x_field)), x_field.detach()[:, 0, 0], axis=0), x_field.detach()[0, :, 1]))


# Planar flow

planar_transform = trans.Planar(input_dim=2)

# Setting the params to fixed values

planar_transform.bias = nn.Parameter(data=torch.tensor([0.5]))
planar_transform.u = nn.Parameter(data=torch.tensor([1.0, 1.0]))
planar_transform.w = nn.Parameter(data=torch.tensor([1.0, 1.0]))


# Computing the pdf with jacobian adjustment 

y_field = planar_transform(x_field)
planar_pdf = torch.exp(base_dist.log_prob(x_field) - planar_transform.log_abs_det_jacobian(x_field,y_field))


# Checking if the transformed distribution integrates to 1

print(np.trapz(np.trapz(planar_pdf.detach(), y_field.detach()[:, 0, 0], axis=0), y_field.detach()[0, :, 1]))
2

I think this is essentially happening because the grid you are evaluating the trapezoidal integral on is irregular (each row has a different set of x-values). This is caused by the transform warping the space and means that the trapezoidal method in 2D isn’t valid…

Evaluating the integral on a regular grid does work. Try this:

transform = T.Planar(input_dim=2).inv
flow_dist = dist.TransformedDistribution(base_dist, [transform1])
planar_pdf = torch.exp(flow_dist.log_prob(x_field))

print(np.trapz(np.trapz(planar_pdf.detach(), x_field.detach()[:, 0, 0], axis=0), x_field.detach()[0, :, 1]))

# plot density
plt.figure(figsize=(5,5))
plt.contourf(x1_s, x2_s, planar_pdf.detach().numpy())
plt.show()

If you want to estimate the normalization constant using the forward method, you can try Monte Carlo integration.

Thanks for the question! This has inspired me to write some more unit tests :slight_smile:

3

Thanks! I didn’t know that ‘.inv’ existed …especially as there is no analytic inverse for this transform and ._inverse doesnt return the inverse transform.

Is there any pyro documentation that talks about ‘.inv’

4

No worries! .inv is inherited from PyTorch’sTransform interface: Probability distributions - torch.distributions — PyTorch 1.13 documentation

I agree, there should be better documentation! I’ve started the first normalizing flows tutorial here: https://nbviewer.jupyter.org/github/stefanwebb/pyro/blob/normalizing-flows-tutorial1/tutorial/source/normalizing_flows_i.ipynb. Hopefully I can discuss features like .inv in a future tutorial