import torch.nn as nn
import torch
from torch.optim.lr_scheduler import ReduceLROnPlateau,OneCycleLR,CyclicLR
import pandas as pd
from sklearn.preprocessing import StandardScaler,MinMaxScaler
import matplotlib.pyplot as plt
from torch.distributions import MultivariateNormal, LogNormal,Normal, Chi2
from torch.distributions.distribution import Distribution
from sklearn.metrics import r2_score
import numpy as np
# It's a distribution that is a kernel density estimate of a Gaussian distribution
class GaussianKDE(Distribution):
def __init__(self, X, bw):
"""
X : tensor (n, d)
`n` points with `d` dimensions to which KDE will be fit
bw : numeric
bandwidth for Gaussian kernel
"""
self.X = X
self.bw = bw
self.dims = X.shape[-1]
self.n = X.shape[0]
self.mvn = MultivariateNormal(loc=torch.zeros(self.dims),
scale_tril=torch.eye(self.dims))
def sample(self, num_samples):
"""
We are sampling from a normal distribution with mean equal to the data points in the dataset and
standard deviation equal to the bandwidth
:param num_samples: the number of samples to draw from the KDE
:return: A sample of size num_samples from the KDE.
"""
idxs = (np.random.uniform(0, 1, num_samples) * self.n).astype(int)
norm = Normal(loc=self.X[idxs], scale=self.bw)
return norm.sample()
def score_samples(self, Y, X=None):
"""Returns the kernel density estimates of each point in `Y`.
Parameters
----------
Y : tensor (m, d)
`m` points with `d` dimensions for which the probability density will
be calculated
X : tensor (n, d), optional
`n` points with `d` dimensions to which KDE will be fit. Provided to
allow batch calculations in `log_prob`. By default, `X` is None and
all points used to initialize KernelDensityEstimator are included.
Returns
-------
log_probs : tensor (m)
log probability densities for each of the queried points in `Y`
"""
if X == None:
X = self.X
log_probs = self.mvn.log_prob((X.unsqueeze(1) - Y)).sum(dim=0)
return log_probs
def log_prob(self, Y):
"""Returns the total log probability of one or more points, `Y`, using
a Multivariate Normal kernel fit to `X` and scaled using `bw`.
Parameters
----------
Y : tensor (m, d)
`m` points with `d` dimensions for which the probability density will
be calculated
Returns
-------
log_prob : numeric
total log probability density for the queried points, `Y`
"""
X_chunks = self.X
Y_chunks = Y
self.Y = Y
log_prob = 0
for x in X_chunks:
for y in Y_chunks:
log_prob += self.score_samples(y,x).sum(dim=0)
return log_prob
class Chi2KDE(Distribution):
def __init__(self, X, bw):
"""
X : tensor (n, d)
`n` points with `d` dimensions to which KDE will be fit
bw : numeric
bandwidth for Gaussian kernel
"""
self.X = X
self.bw = bw
self.dims = X.shape[-1]
self.n = X.shape[0]
self.mvn = Chi2(self.dims)
def sample(self, num_samples):
idxs = (np.random.uniform(0, 1, num_samples) * self.n).astype(int)
norm = LogNormal(loc=self.X[idxs], scale=self.bw)
return norm.sample()
def score_samples(self, Y, X=None):
"""Returns the kernel density estimates of each point in `Y`.
Parameters
----------
Y : tensor (m, d)
`m` points with `d` dimensions for which the probability density will
be calculated
X : tensor (n, d), optional
`n` points with `d` dimensions to which KDE will be fit. Provided to
allow batch calculations in `log_prob`. By default, `X` is None and
all points used to initialize KernelDensityEstimator are included.
Returns
-------
log_probs : tensor (m)
log probability densities for each of the queried points in `Y`
"""
if X == None:
X = self.X
log_probs = self.mvn.log_prob(abs(X.unsqueeze(1) - Y)).sum()
return log_probs
def log_prob(self, Y):
"""Returns the total log probability of one or more points, `Y`, using
a Multivariate Normal kernel fit to `X` and scaled using `bw`.
Parameters
----------
Y : tensor (m, d)
`m` points with `d` dimensions for which the probability density will
be calculated
Returns
-------
log_prob : numeric
total log probability density for the queried points, `Y`
"""
X_chunks = self.X
Y_chunks = Y
self.Y = Y
log_prob = 0
for x in X_chunks:
for y in Y_chunks:
log_prob += self.score_samples(y,x).sum(dim=0)
return log_prob
class PlanarFlow(nn.Module):
"""
A single planar flow, computes T(x) and log(det(jac_T)))
"""
def __init__(self, D):
super(PlanarFlow, self).__init__()
self.u = nn.Parameter(torch.Tensor(1, D), requires_grad=True)
self.w = nn.Parameter(torch.Tensor(1, D), requires_grad=True)
self.b = nn.Parameter(torch.Tensor(1), requires_grad=True)
self.h = torch.tanh
self.init_params()
def init_params(self):
self.w.data.uniform_(0.4, 1)
self.b.data.uniform_(0.4, 1)
self.u.data.uniform_(0.4, 1)
def forward(self, z):
linear_term = torch.mm(z, self.w.T) + self.b
return z + self.u * self.h(linear_term)
def h_prime(self, x):
"""
Derivative of tanh
"""
return (1 - self.h(x) ** 2)
def psi(self, z):
inner = torch.mm(z, self.w.T) + self.b
return self.h_prime(inner) * self.w
def log_det(self, z):
inner = 1 + torch.mm(self.psi(z), self.u.T)
return torch.log(torch.abs(inner))
# It's a normalizing flow that takes in a distribution and outputs a distribution.
class NormalizingFlow(nn.Module):
"""
A normalizng flow composed of a sequence of planar flows.
"""
def __init__(self, D, n_flows=2):
"""
The function takes in two arguments, D and n_flows. D is the dimension of the data, and n_flows
is the number of flows. The function then creates a list of PlanarFlow objects, where the number
of PlanarFlow objects is equal to n_flows
:param D: the dimensionality of the data
:param n_flows: number of flows to use, defaults to 2 (optional)
"""
super(NormalizingFlow, self).__init__()
self.flows = nn.ModuleList(
[PlanarFlow(D) for _ in range(n_flows)])
def sample(self, base_samples):
"""
Transform samples from a simple base distribution
by passing them through a sequence of Planar flows.
"""
samples = base_samples
for flow in self.flows:
samples = flow(samples)
return samples
def forward(self, x):
"""
Computes and returns the sum of log_det_jacobians
and the transformed samples T(x).
"""
sum_log_det = 0
transformed_sample = x
for i in range(len(self.flows)):
log_det_i = (self.flows[i].log_det(transformed_sample))
sum_log_det += log_det_i
transformed_sample = self.flows[i](transformed_sample)
return transformed_sample, sum_log_det
def random_normal_samples(n, dim=2):
return torch.zeros(n, dim).normal_(mean=0, std=1.5)
class nflow():
def __init__(self,dim=2,latent=16,batchsize:int=1,dataset=None):
"""
The function __init__ initializes the class NormalizingFlowModel with the parameters dim,
latent, batchsize, and datasetPath
:param dim: The dimension of the data, defaults to 2 (optional)
:param latent: The number of latent variables in the model, defaults to 16 (optional)
:param batchsize: The number of samples to generate at a time, defaults to 1
:type batchsize: int (optional)
:param datasetPath: The path to the dataset, defaults to data/dataset.csv
:type datasetPath: str (optional)
"""
self.dim = dim
self.batchsize = batchsize
self.model = NormalizingFlow(dim, latent)
self.dataset = dataset
def compile(self,optim:torch.optim=torch.optim.Adam,distribution:str='GaussianKDE',lr:float=0.00015,bw:float=0.1,wd=0.0015):
"""
It takes in a dataset, a model, and a distribution, and returns a compiled model
:param optim: the optimizer to use
:type optim: torch.optim
:param distribution: the type of distribution to use, defaults to GaussianKDE
:type distribution: str (optional)
:param lr: learning rate
:type lr: float
:param bw: bandwidth for the KDE
:type bw: float
"""
if wd:
self.opt = optim(
params=self.model.parameters(),
lr=lr,
weight_decay = wd
# momentum=0.9
# momentum=0.1
)
else:
self.opt = optim(
params=self.model.parameters(),
lr=lr,
# momentum=0.9
# momentum=0.1
)
self.scaler = StandardScaler()
self.scaler_mm = MinMaxScaler(feature_range=(0,1))
df = pd.read_csv(self.dataset)
df = df.iloc[:,1:]
if 'Chi2' in distribution:
self.scaled=self.scaler_mm.fit_transform(df)
else: self.scaled = self.scaler.fit_transform(df)
self.density = globals()[distribution](X=torch.tensor(self.scaled, dtype=torch.float32), bw=bw)
# self.dl = torch.utils.data.DataLoader(scaled,batchsize=self.batchsize)
self.scheduler = ReduceLROnPlateau(self.opt, patience=10000)
self.losses = []
def train(self,iters:int=1000):
"""
> We sample from a normal distribution, pass the samples through the model, and then calculate
the loss
:param iters: number of iterations to train for, defaults to 1000
:type iters: int (optional)
"""
for idx in range(iters):
if idx % 100 == 0:
print("Iteration {}".format(idx))
samples = torch.autograd.Variable(random_normal_samples(self.batchsize,self.dim))
z_k, sum_log_det = self.model(samples)
log_p_x = self.density.log_prob(z_k)
# Reverse KL since we can evaluate target density but can't sample
loss = (-sum_log_det - (log_p_x)).mean()
self.opt.zero_grad()
loss.backward()
self.opt.step()
self.scheduler.step(loss)
self.losses.append(loss.item())
if idx % 100 == 0:
print("Loss {}".format(loss.item()))
yield idx,loss.item()
def performance(self):
"""
The function takes the model and the scaled data as inputs, samples from the model, and then
prints the r2 score of the samples and the scaled data.
"""
samples = ((self.model.sample(torch.tensor(self.scaled).float())).detach().numpy())
print('r2', r2_score(self.scaled,samples))