class LatentDistanceGaussianWeightDistribution(GaussianWeightDistribution, GibbsSampling):
"""
l_n ~ N(0, sigma^2 I)
W_{n', n} ~ N(A * ||l_{n'} - l_{n}||_2^2 + b, ) for n' != n
"""
def __init__(self, N, B=1, dim=2,
b=0.5,
sigma=None, Sigma_0=None, nu_0=None,
mu_self=0.0, eta=0.01):
"""
Initialize SBM with parameters defined above.
"""
super(LatentDistanceGaussianWeightDistribution, self).__init__(N)
self.B = B
self.dim = dim
self.b = b
self.eta = eta
self.L = np.sqrt(eta) * np.random.randn(N,dim)
if Sigma_0 is None:
Sigma_0 = np.eye(B)
if nu_0 is None:
nu_0 = B + 2
self.cov = GaussianFixedMean(mu=np.zeros(B), sigma=sigma, lmbda_0=Sigma_0, nu_0=nu_0)
# Special case self-weights (along the diagonal)
self._self_gaussian = Gaussian(mu_0=mu_self*np.ones(B),
sigma_0=Sigma_0,
nu_0=nu_0,
kappa_0=1.0)
@property
def D(self):
# return np.sqrt(((self.L[:, None, :] - self.L[None, :, :]) ** 2).sum(2))
return ((self.L[:, None, :] - self.L[None, :, :]) ** 2).sum(2)
@property
def Mu(self):
Mu = -self.D + self.b
Mu = np.tile(Mu[:,:,None], (1,1,self.B))
for n in xrange(self.N):
Mu[n,n,:] = self._self_gaussian.mu
return Mu
@property
def Sigma(self):
sig = self.cov.sigma
Sig = np.tile(sig[None,None,:,:], (self.N, self.N,1,1))
for n in xrange(self.N):
Sig[n,n,:,:] = self._self_gaussian.sigma
return Sig
def initialize_from_prior(self):
self.L = np.sqrt(self.eta) * np.random.randn(self.N, self.dim)
self.cov.resample()
def initialize_hypers(self, W):
# Optimize the initial locations
self._optimize_L(np.ones((self.N,self.N)), W)
def log_prior(self):
"""
Compute the prior probability of F, mu0, and lmbda
"""
from graphistician.internals.utils import \
normal_inverse_wishart_log_prob, \
inverse_wishart_log_prob
lp = 0
# Log prior of F under spherical Gaussian prior
from scipy.stats import norm, invgamma
lp += invgamma.logpdf(self.eta, 1, 1)
lp += norm.logpdf(self.b, 0, 1)
lp += norm.logpdf(self.L, 0, 1).sum()
lp += inverse_wishart_log_prob(self.cov)
lp += normal_inverse_wishart_log_prob(self._self_gaussian)
# Log prior of mu_0 and mu_self
return lp
def _hmc_log_probability(self, L, b, A, W):
"""
Compute the log probability as a function of L.
This allows us to take the gradients wrt L using autograd.
:param L:
:param A:
:return:
"""
assert self.B == 1
import autograd.numpy as anp
# Compute pairwise distance
L1 = anp.reshape(L,(self.N,1,self.dim))
L2 = anp.reshape(L,(1,self.N,self.dim))
# Mu = a * anp.sqrt(anp.sum((L1-L2)**2, axis=2)) + b
Mu = -anp.sum((L1-L2)**2, axis=2) + b
Aoff = A * (1-anp.eye(self.N))
X = (W - Mu[:,:,None]) * Aoff[:,:,None]
# Get the covariance and precision
Sig = self.cov.sigma[0,0]
Lmb = 1./Sig
lp = anp.sum(-0.5 * X**2 * Lmb)
# Log prior of L under spherical Gaussian prior
lp += -0.5 * anp.sum(L * L / self.eta)
# Log prior of mu0 under standardGaussian prior
lp += -0.5 * b ** 2
return lp
def sample_predictive_parameters(self):
Lext = \
np.vstack((self.L, np.sqrt(self.eta) * np.random.randn(1, self.dim)))
# Compute mean and covariance over extended space
D = ((Lext[:,None,:] - Lext[None,:,:])**2).sum(2)
Mu = -D + self.b
Mu_row = np.tile(Mu[-1,:][:,None], (1,self.B))
Mu_row[-1] = self._self_gaussian.mu
Mu_col = Mu_row.copy()
# Mu = np.tile(Mu[:,:,None], (1,1,self.B))
# for n in xrange(self.N+1):
# Mu[n,n,:] = self._self_gaussian.mu
L = np.linalg.cholesky(self.cov.sigma)
L_row = np.tile(L[None,:,:], (self.N+1, 1, 1))
L_row[-1] = np.linalg.cholesky(self._self_gaussian.sigma)
L_col = L_row.copy()
# L = np.tile(L[None,None,:,:], (self.N+1, self.N+1, 1, 1))
# for n in xrange(self.N+1):
# L[n,n,:,:] = np.linalg.cholesky(self._self_gaussian.sigma)
# Mu_row, Mu_col = Mu[-1,:,:], Mu[:,-1,:]
# L_row, L_col = L[-1,:,:,:], L[:,-1,:,:]
return Mu_row, Mu_col, L_row, L_col
def resample(self, (A,W)):
self._resample_L(A, W)
self._resample_b(A, W)
self._resample_cov(A, W)
self._resample_self_gaussian(A, W)
self._resample_eta()
class SBMGaussianWeightSharedCov(SBMGaussianWeightDistribution):
def __init__(self, N, B=1,
C=3, pi=10.0,
mu_0=None, Sigma_0=None, nu_0=None,
special_case_self_conns=True):
super(SBMGaussianWeightSharedCov, self).\
__init__(N, B=B, C=C, pi=pi,
mu_0=mu_0, Sigma_0=Sigma_0, nu_0=nu_0,
special_case_self_conns=special_case_self_conns)
if mu_0 is None:
mu_0 = np.zeros(B)
if Sigma_0 is None:
Sigma_0 = np.eye(B)
if nu_0 is None:
nu_0 = B + 2
self._cov_model = GaussianFixedMean(mu=np.zeros(B),
nu_0=nu_0, lmbda_0=Sigma_0)
self._gaussians = [[GaussianFixedCov(mu_0=mu_0, sigma_0=np.eye(B),
sigma=self._cov_model.sigma)
for _ in xrange(C)]
for _ in xrange(C)]
def resample_mu_and_Sig(self, A, W):
"""
Resample p given observations of the weights
"""
Abool = A.astype(np.bool)
for c1 in xrange(self.C):
for c2 in xrange(self.C):
mask = self._get_mask(Abool, c1, c2)
self._gaussians[c1][c2].resample(W[mask])
# Resample self connection
if self.special_case_self_conns:
mask = np.eye(self.N, dtype=np.bool) & Abool
self._self_gaussian.resample(W[mask])
# Resample covariance
A_offdiag = Abool.copy()
np.fill_diagonal(A_offdiag, False)
W_cent = (W - self.Mu)[A_offdiag]
self._cov_model.resample(W_cent)
# Update gaussians
for c1 in xrange(self.C):
for c2 in xrange(self.C):
self._gaussians[c1][c2].sigma = self._cov_model.sigma
def log_prior(self):
"""
Compute the log likelihood of a set of SBM parameters
:param x: (m,p,v) tuple
:return:
"""
from scipy.stats import dirichlet
lp = 0
# Get the log probability of the block probabilities
lp += dirichlet(self.pi).logpdf(self.m)
# Get the prior probability of the Gaussian parameters under NIW prior
# for c1 in xrange(self.C):
# for c2 in xrange(self.C):
# lp += normal_inverse_wishart_log_prob(self._gaussians[c1][c2])
#
# if self.special_case_self_conns:
# lp += normal_inverse_wishart_log_prob(self._self_gaussian)
# Get the probability of the block assignments
lp += (np.log(self.m)[self.c]).sum()
return lp
def initialize_hypers(self, W):
mu_0 = W.mean(axis=(0,1))
sigma_0 = np.diag(W.var(axis=(0,1)))
# Set the global cov
nu_0 = self._cov_model.nu_0
self._cov_model.sigma_0 = sigma_0 * (nu_0 - self.B - 1)
# Set the mean
for c1 in xrange(self.C):
for c2 in xrange(self.C):
self._gaussians[c1][c2].mu_0 = mu_0
self._gaussians[c1][c2].sigma = self._cov_model.sigma_0
self._gaussians[c1][c2].resample()
if self.special_case_self_conns:
W_self = W[np.arange(self.N), np.arange(self.N)]
self._self_gaussian.mu_0 = W_self.mean(axis=0)
self._self_gaussian.sigma_0 = np.diag(W_self.var(axis=0))
self._self_gaussian.resample()
# Cluster the neurons based on their rows and columns
from sklearn.cluster import KMeans
features = np.hstack((W[:,:,0], W[:,:,0].T))
km = KMeans(n_clusters=self.C)
km.fit(features)
self.c = km.labels_.astype(np.int)
print "Initial c: ", self.c
class SBMGaussianWeightSharedCov(SBMGaussianWeightDistribution):
def __init__(self,
N,
B=1,
C=3,
pi=10.0,
mu_0=None,
Sigma_0=None,
nu_0=None,
special_case_self_conns=True):
super(SBMGaussianWeightSharedCov, self).\
__init__(N, B=B, C=C, pi=pi,
mu_0=mu_0, Sigma_0=Sigma_0, nu_0=nu_0,
special_case_self_conns=special_case_self_conns)
if mu_0 is None:
mu_0 = np.zeros(B)
if Sigma_0 is None:
Sigma_0 = np.eye(B)
if nu_0 is None:
nu_0 = B + 2
self._cov_model = GaussianFixedMean(mu=np.zeros(B),
nu_0=nu_0,
lmbda_0=Sigma_0)
self._gaussians = [[
GaussianFixedCov(mu_0=mu_0,
sigma_0=np.eye(B),
sigma=self._cov_model.sigma) for _ in xrange(C)
] for _ in xrange(C)]
def resample_mu_and_Sig(self, A, W):
"""
Resample p given observations of the weights
"""
Abool = A.astype(np.bool)
for c1 in xrange(self.C):
for c2 in xrange(self.C):
mask = self._get_mask(Abool, c1, c2)
self._gaussians[c1][c2].resample(W[mask])
# Resample self connection
if self.special_case_self_conns:
mask = np.eye(self.N, dtype=np.bool) & Abool
self._self_gaussian.resample(W[mask])
# Resample covariance
A_offdiag = Abool.copy()
np.fill_diagonal(A_offdiag, False)
W_cent = (W - self.Mu)[A_offdiag]
self._cov_model.resample(W_cent)
# Update gaussians
for c1 in xrange(self.C):
for c2 in xrange(self.C):
self._gaussians[c1][c2].sigma = self._cov_model.sigma
def log_prior(self):
"""
Compute the log likelihood of a set of SBM parameters
:param x: (m,p,v) tuple
:return:
"""
from scipy.stats import dirichlet
lp = 0
# Get the log probability of the block probabilities
lp += dirichlet(self.pi).logpdf(self.m)
# Get the prior probability of the Gaussian parameters under NIW prior
# for c1 in xrange(self.C):
# for c2 in xrange(self.C):
# lp += normal_inverse_wishart_log_prob(self._gaussians[c1][c2])
#
# if self.special_case_self_conns:
# lp += normal_inverse_wishart_log_prob(self._self_gaussian)
# Get the probability of the block assignments
lp += (np.log(self.m)[self.c]).sum()
return lp
def initialize_hypers(self, W):
mu_0 = W.mean(axis=(0, 1))
sigma_0 = np.diag(W.var(axis=(0, 1)))
# Set the global cov
nu_0 = self._cov_model.nu_0
self._cov_model.sigma_0 = sigma_0 * (nu_0 - self.B - 1)
# Set the mean
for c1 in xrange(self.C):
for c2 in xrange(self.C):
self._gaussians[c1][c2].mu_0 = mu_0
self._gaussians[c1][c2].sigma = self._cov_model.sigma_0
self._gaussians[c1][c2].resample()
if self.special_case_self_conns:
W_self = W[np.arange(self.N), np.arange(self.N)]
self._self_gaussian.mu_0 = W_self.mean(axis=0)
self._self_gaussian.sigma_0 = np.diag(W_self.var(axis=0))
self._self_gaussian.resample()
# Cluster the neurons based on their rows and columns
from sklearn.cluster import KMeans
features = np.hstack((W[:, :, 0], W[:, :, 0].T))
km = KMeans(n_clusters=self.C)
km.fit(features)
self.c = km.labels_.astype(np.int)
print "Initial c: ", self.c
class LatentDistanceGaussianWeightDistribution(GaussianWeightDistribution,
GibbsSampling):
"""
l_n ~ N(0, sigma^2 I)
W_{n', n} ~ N(A * ||l_{n'} - l_{n}||_2^2 + b, ) for n' != n
"""
def __init__(self,
N,
B=1,
dim=2,
b=0.5,
sigma=None,
Sigma_0=None,
nu_0=None,
mu_self=0.0,
eta=0.01):
"""
Initialize SBM with parameters defined above.
"""
super(LatentDistanceGaussianWeightDistribution, self).__init__(N)
self.B = B
self.dim = dim
self.b = b
self.eta = eta
self.L = np.sqrt(eta) * np.random.randn(N, dim)
if Sigma_0 is None:
Sigma_0 = np.eye(B)
if nu_0 is None:
nu_0 = B + 2
self.cov = GaussianFixedMean(mu=np.zeros(B),
sigma=sigma,
lmbda_0=Sigma_0,
nu_0=nu_0)
# Special case self-weights (along the diagonal)
self._self_gaussian = Gaussian(mu_0=mu_self * np.ones(B),
sigma_0=Sigma_0,
nu_0=nu_0,
kappa_0=1.0)
@property
def D(self):
# return np.sqrt(((self.L[:, None, :] - self.L[None, :, :]) ** 2).sum(2))
return ((self.L[:, None, :] - self.L[None, :, :])**2).sum(2)
@property
def Mu(self):
Mu = -self.D + self.b
Mu = np.tile(Mu[:, :, None], (1, 1, self.B))
for n in xrange(self.N):
Mu[n, n, :] = self._self_gaussian.mu
return Mu
@property
def Sigma(self):
sig = self.cov.sigma
Sig = np.tile(sig[None, None, :, :], (self.N, self.N, 1, 1))
for n in xrange(self.N):
Sig[n, n, :, :] = self._self_gaussian.sigma
return Sig
def initialize_from_prior(self):
self.L = np.sqrt(self.eta) * np.random.randn(self.N, self.dim)
self.cov.resample()
def initialize_hypers(self, W):
# Optimize the initial locations
self._optimize_L(np.ones((self.N, self.N)), W)
def log_prior(self):
"""
Compute the prior probability of F, mu0, and lmbda
"""
from graphistician.internals.utils import \
normal_inverse_wishart_log_prob, \
inverse_wishart_log_prob
lp = 0
# Log prior of F under spherical Gaussian prior
from scipy.stats import norm, invgamma
lp += invgamma.logpdf(self.eta, 1, 1)
lp += norm.logpdf(self.b, 0, 1)
lp += norm.logpdf(self.L, 0, 1).sum()
lp += inverse_wishart_log_prob(self.cov)
lp += normal_inverse_wishart_log_prob(self._self_gaussian)
# Log prior of mu_0 and mu_self
return lp
def _hmc_log_probability(self, L, b, A, W):
"""
Compute the log probability as a function of L.
This allows us to take the gradients wrt L using autograd.
:param L:
:param A:
:return:
"""
assert self.B == 1
import autograd.numpy as anp
# Compute pairwise distance
L1 = anp.reshape(L, (self.N, 1, self.dim))
L2 = anp.reshape(L, (1, self.N, self.dim))
# Mu = a * anp.sqrt(anp.sum((L1-L2)**2, axis=2)) + b
Mu = -anp.sum((L1 - L2)**2, axis=2) + b
Aoff = A * (1 - anp.eye(self.N))
X = (W - Mu[:, :, None]) * Aoff[:, :, None]
# Get the covariance and precision
Sig = self.cov.sigma[0, 0]
Lmb = 1. / Sig
lp = anp.sum(-0.5 * X**2 * Lmb)
# Log prior of L under spherical Gaussian prior
lp += -0.5 * anp.sum(L * L / self.eta)
# Log prior of mu0 under standardGaussian prior
lp += -0.5 * b**2
return lp
def sample_predictive_parameters(self):
Lext = \
np.vstack((self.L, np.sqrt(self.eta) * np.random.randn(1, self.dim)))
# Compute mean and covariance over extended space
D = ((Lext[:, None, :] - Lext[None, :, :])**2).sum(2)
Mu = -D + self.b
Mu_row = np.tile(Mu[-1, :][:, None], (1, self.B))
Mu_row[-1] = self._self_gaussian.mu
Mu_col = Mu_row.copy()
# Mu = np.tile(Mu[:,:,None], (1,1,self.B))
# for n in xrange(self.N+1):
# Mu[n,n,:] = self._self_gaussian.mu
L = np.linalg.cholesky(self.cov.sigma)
L_row = np.tile(L[None, :, :], (self.N + 1, 1, 1))
L_row[-1] = np.linalg.cholesky(self._self_gaussian.sigma)
L_col = L_row.copy()
# L = np.tile(L[None,None,:,:], (self.N+1, self.N+1, 1, 1))
# for n in xrange(self.N+1):
# L[n,n,:,:] = np.linalg.cholesky(self._self_gaussian.sigma)
# Mu_row, Mu_col = Mu[-1,:,:], Mu[:,-1,:]
# L_row, L_col = L[-1,:,:,:], L[:,-1,:,:]
return Mu_row, Mu_col, L_row, L_col
def resample(self, (A, W)):
self._resample_L(A, W)
self._resample_b(A, W)
self._resample_cov(A, W)
self._resample_self_gaussian(A, W)
self._resample_eta()
class _LatentDistanceModelGaussianMixin(_NetworkModel):
"""
l_n ~ N(0, sigma^2 I)
W_{n', n} ~ N(A * ||l_{n'} - l_{n}||_2^2 + b, ) for n' != n
"""
def __init__(self,
N,
B=1,
dim=2,
b=0.5,
sigma=None,
Sigma_0=None,
nu_0=None,
mu_self=0.0,
eta=0.01):
super(_LatentDistanceModelGaussianMixin, self).__init__(N, B)
self.B = B
self.dim = dim
self.b = b
self.eta = eta
self.L = np.sqrt(eta) * np.random.randn(N, dim)
if Sigma_0 is None:
Sigma_0 = np.eye(B)
if nu_0 is None:
nu_0 = B + 2
self.cov = GaussianFixedMean(mu=np.zeros(B),
sigma=sigma,
lmbda_0=Sigma_0,
nu_0=nu_0)
# Special case self-weights (along the diagonal)
self._self_gaussian = Gaussian(mu_0=mu_self * np.ones(B),
sigma_0=Sigma_0,
nu_0=nu_0,
kappa_0=1.0)
@property
def D(self):
# return np.sqrt(((self.L[:, None, :] - self.L[None, :, :]) ** 2).sum(2))
return ((self.L[:, None, :] - self.L[None, :, :])**2).sum(2)
@property
def mu_W(self):
Mu = -self.D + self.b
Mu = np.tile(Mu[:, :, None], (1, 1, self.B))
for n in range(self.N):
Mu[n, n, :] = self._self_gaussian.mu
return Mu
@property
def sigma_W(self):
sig = self.cov.sigma
Sig = np.tile(sig[None, None, :, :], (self.N, self.N, 1, 1))
for n in range(self.N):
Sig[n, n, :, :] = self._self_gaussian.sigma
return Sig
def initialize_from_prior(self):
self.L = np.sqrt(self.eta) * np.random.randn(self.N, self.dim)
self.cov.resample()
def initialize_hypers(self, W):
# Optimize the initial locations
self._optimize_L(np.ones((self.N, self.N)), W)
def _hmc_log_probability(self, L, b, A, W):
"""
Compute the log probability as a function of L.
This allows us to take the gradients wrt L using autograd.
:param L:
:param A:
:return:
"""
assert self.B == 1
import autograd.numpy as atnp
# Compute pairwise distance
L1 = atnp.reshape(L, (self.N, 1, self.dim))
L2 = atnp.reshape(L, (1, self.N, self.dim))
# Mu = a * anp.sqrt(anp.sum((L1-L2)**2, axis=2)) + b
Mu = -atnp.sum((L1 - L2)**2, axis=2) + b
Aoff = A * (1 - atnp.eye(self.N))
X = (W - Mu[:, :, None]) * Aoff[:, :, None]
# Get the covariance and precision
Sig = self.cov.sigma[0, 0]
Lmb = 1. / Sig
lp = atnp.sum(-0.5 * X**2 * Lmb)
# Log prior of L under spherical Gaussian prior
lp += -0.5 * atnp.sum(L * L / self.eta)
# Log prior of mu0 under standardGaussian prior
lp += -0.5 * b**2
return lp
def resample(self, data=[]):
super(_LatentDistanceModelGaussianMixin, self).resample(data)
A, W = data
N, B = self.N, self.B
self._resample_L(A, W)
self._resample_b(A, W)
self._resample_cov(A, W)
self._resample_self_gaussian(A, W)
self._resample_eta()
# print "eta: ", self.eta, "\tb: ", self.b
def _resample_L(self, A, W):
"""
Resample the locations given A
:return:
"""
from autograd import grad
from hips.inference.hmc import hmc
lp = lambda L: self._hmc_log_probability(L, self.b, A, W)
dlp = grad(lp)
stepsz = 0.005
nsteps = 10
# lp0 = lp(self.L)
self.L = hmc(lp,
dlp,
stepsz,
nsteps,
self.L.copy(),
negative_log_prob=False)
# lpf = lp(self.L)
# print "diff lp: ", (lpf - lp0)
def _optimize_L(self, A, W):
"""
Resample the locations given A
:return:
"""
import autograd.numpy as atnp
from autograd import grad
from scipy.optimize import minimize
lp = lambda Lflat: -self._hmc_log_probability(
atnp.reshape(Lflat, (self.N, 2)), self.b, A, W)
dlp = grad(lp)
res = minimize(lp, np.ravel(self.L), jac=dlp, method="bfgs")
self.L = np.reshape(res.x, (self.N, 2))
def _resample_b_hmc(self, A, W):
"""
Resample the distance dependence offset
:return:
"""
# TODO: We could sample from the exact Gaussian conditional
from autograd import grad
from hips.inference.hmc import hmc
lp = lambda b: self._hmc_log_probability(self.L, b, A, W)
dlp = grad(lp)
stepsz = 0.0001
nsteps = 10
b = hmc(lp,
dlp,
stepsz,
nsteps,
np.array(self.b),
negative_log_prob=False)
self.b = float(b)
print("b: ", self.b)
def _resample_b(self, A, W):
"""
Resample the distance dependence offset
W ~ N(mu, sigma)
= N(-D + b, sigma)
implies
W + D ~ N(b, sigma).
If b ~ N(0, 1), we can compute the Gaussian conditional
in closed form.
"""
D = self.D
sigma = self.cov.sigma[0, 0]
Aoff = (A * (1 - np.eye(self.N))).astype(np.bool)
X = (W + D[:, :, None])[Aoff]
# Now X ~ N(b, sigma)
mu0, sigma0 = 0.0, 1.0
N = X.size
sigma_post = 1. / (1. / sigma0 + N / sigma)
mu_post = sigma_post * (mu0 / sigma0 + X.sum() / sigma)
self.b = mu_post + np.sqrt(sigma_post) * np.random.randn()
# print "b: ", self.b
def _resample_cov(self, A, W):
# Resample covariance matrix
Mu = self.Mu
mask = (True - np.eye(self.N, dtype=np.bool)) & A.astype(np.bool)
self.cov.resample(W[mask] - Mu[mask])
def _resample_self_gaussian(self, A, W):
# Resample self connection
mask = np.eye(self.N, dtype=np.bool) & A.astype(np.bool)
self._self_gaussian.resample(W[mask])
def _resample_eta(self):
"""
Resample sigma under an inverse gamma prior, sigma ~ IG(1,1)
:return:
"""
L = self.L
a_prior = 1.0
b_prior = 1.0
a_post = a_prior + L.size / 2.0
b_post = b_prior + (L**2).sum() / 2.0
from scipy.stats import invgamma
self.eta = invgamma.rvs(a=a_post, scale=b_post)