def __init__(self, positions, heterogeneity, linear_predictor):
super().__init__()
self._deterministic = True
N, K = positions.shape()
# neighboring nodes
self.positions = positions
self.heterogeneity = heterogeneity
self.linear_predictor = linear_predictor
# hidden nodes
self._products = GaussianArray.uniform((N, N, K))
self._heterogeneity_expanded = GaussianArray.uniform((N, N, 2))
self._vector = GaussianArray.uniform((N, N, K + 2))
self._nodes.update({
"products": self._products,
"heterogeneity_expanded": self._heterogeneity_expanded,
"vector": self._vector
})
# hidden factors
self._product = Product(parent=self.positions, child=self._products)
self._expand_transpose = ExpandTranspose(
child=self._heterogeneity_expanded, parent=self.heterogeneity)
self._concatenate = Concatenate(parts={
"products":
self._products,
"heterogeneity_expanded":
self._heterogeneity_expanded
},
vector=self._vector)
self._sum = Sum(parent=self._vector, child=self.linear_predictor)
self._factors.update({
"product": self._product,
"expand_transpose": self._expand_transpose,
"concatenate": self._concatenate,
"sum": self._sum
})
def __init__(self, child, parent):
super().__init__()
self._deterministic = True
self.child = child
self.parent = parent
self.message_to_child = GaussianArray.uniform(child.shape())
self.message_to_parent = GaussianArray.uniform(parent.shape())
def __init__(self, shape_in, shape_out):
super().__init__()
d = len(shape_out)
self.message_to_x = {k: GaussianArray.uniform(s) for k, s in shape_in.items()}
self.message_to_v = GaussianArray.uniform(shape_out)
size = [s[-1] for k, s in shape_in.items()]
begin = [0, *np.cumsum(size[:-1])]
self._size = [tuple([*shape_out[:-1], s]) for s in size]
self._begin = [tuple([*[0]*(d - 1), s]) for s in begin]
self._name = [k for k, s in shape_in.items()]
def __init__(self, child: GaussianArray, parent: GaussianArray, variance: float = 1.):
super().__init__()
self._deterministic = False
self.shape = child.shape()
self.child = child
self.parent = parent
self.message_to_child = GaussianArray.uniform(self.shape)
self.message_to_parent = GaussianArray.uniform(self.shape)
self.variance = ParameterArrayLogScale(variance, False, name="AddVariance.variance")
self._parameters = {"variance": self.variance}
def __init__(self, child: GaussianArray, mean: float = 0., variance: float = 1., initial=None, name=""):
super().__init__()
self._deterministic = False
self.mean = ParameterArray(mean, True, name=name+".mean")
self.variance = ParameterArrayLogScale(variance, True, name=name+".variance")
self._parameters = {"mean": self.mean, "variance": self.variance}
self.child = child
self.shape = child.shape()
self.prior = GaussianArray.from_shape(self.shape, self.mean.value(), self.variance.value())
# initialize child
self.message_to_child = GaussianArray.from_array(initial, tf.ones_like(initial) * variance * 0.1)
self.child.set_to(self.message_to_child)
def forward(self):
# products
# if self._current_iter == 0:
# # break symmetry
# # self.factors["product"].message_to_product = GaussianArray.from_array(
# # mean=tf.random.normal((self.N, self.K), 0., 1.),
# # variance=tf.ones((self.N, self.K)) * 10.
# # ) * self.factors["latent_prior"].message_to_x
# # self.nodes["product"] = self.factors["product"].message_to_product * \
# # self.factors["concatenate"].message_to_x["s_uv"]
# # self.factors["product"].message_to_x = GaussianArray.from_array(
# # mean=tf.random.normal((self.N, self.K), 0., 1.),
# # variance=tf.ones((self.N, self.K)) * 10.
# # )
# self.factors["latent_prior"].message_to_x = GaussianArray.from_array(
# mean=tf.random.normal((self.N, self.K), 0., 1.),
# variance=tf.ones((self.N, self.K)) * 1000.
# )
# latent position
self.nodes["latent"] = self.factors["latent_prior"].to_x() * self.factors["product"].message_to_x
self.nodes["product"] = self.factors["product"].to_product(
x=self.nodes["latent"]
) * self.factors["concatenate"].message_to_x["s_uv"]
# heterogeneity
to_alpha = \
self.factors["concatenate"].message_to_x["a_u"].product(0) * \
self.factors["concatenate"].message_to_x["a_v"].product(1)
self.nodes["heterogeneity"] = self.factors["heterogeneity_prior"].message_to_x * to_alpha
# concatenate
alpha = self.nodes["heterogeneity"]
x = {
"a_u": GaussianArray(
tf.tile(tf.expand_dims(alpha.precision(), 0), [self.N, 1, 1]),
tf.tile(tf.expand_dims(alpha.mean_times_precision(), 0), [self.N, 1, 1])
),
"a_v": GaussianArray(
tf.tile(tf.expand_dims(alpha.precision(), 1), [1, self.N, 1]),
tf.tile(tf.expand_dims(alpha.mean_times_precision(), 1), [1, self.N, 1])
),
"s_uv": self.nodes["product"]
}
self.nodes["vector"] = self.factors["concatenate"].to_v(x) * \
self.factors["sum"].to_x(self.nodes["vector"], self.nodes["linear_predictor"])
# linear predictor
self.nodes["linear_predictor"] = \
self.factors["sum"].to_sum(self.nodes["vector"]) * \
self.factors["noise"].message_to_mean
# noizy linear predictor
self.nodes["noisy_linear_predictor"] = self.factors["noise"].to_x(
mean=self.nodes["linear_predictor"],
variance=self.parameters["noise"]._value
) * self.factors["adjacency"].message_to_x
def __init__(self, N, K, A):
self.A = A
self.N = N
self.K = K
self.parameters = {
"noise": ParameterArray(1. * tf.ones((1, 1)))
}
self.nodes = {
"latent": GaussianArray.uniform((N, K)),
"heterogeneity": GaussianArray.uniform((N, 1)),
"product": GaussianArray.uniform((N, N, K)),
"vector": GaussianArray.uniform((N, N, K+2)),
"linear_predictor": GaussianArray.uniform((N, N)),
"noisy_linear_predictor": GaussianArray.uniform((N, N)),
"links": tf.zeros((N, N))
}
self.factors = {
"latent_prior": Prior(GaussianArray.from_shape((N, K), 0., 1.)),
"heterogeneity_prior": Prior(GaussianArray.from_shape((N, 1), 0., 1.)),
"product": Product((N, K), (N, N, K)),
"concatenate": Concatenate({"a_u": (N, N, 1), "a_v": (N, N, 1), "s_uv": (N, N, K)}, (N, N, K+2)),
"sum": Sum((N, N, K+2), (N, N)),
"noise": AddVariance((N, N)),
"adjacency": Probit((N, N))
}
self._current_iter = 0
self._break_symmetry()
def __init__(self,
parent: GaussianArray,
child_cts: GaussianArray,
child_bin: BernoulliArray,
variance_cts=None,
variance_bin=None,
bin_model="NoisyProbit"):
super().__init__()
self._deterministic = False
# nodes
self.parent = parent
self.child_cts = child_cts
self.child_bin = child_bin
# dimensions
N, K = parent.shape()
p_cts, p_bin = child_cts.shape()[1], child_bin.shape()[1]
if variance_cts is None:
variance_cts = tf.ones((1, p_cts))
if variance_bin is None:
variance_bin = tf.ones((1, p_bin))
# hidden nodes
self._lin_pred_cts = GaussianArray.uniform((N, p_cts))
self._lin_pred_bin = GaussianArray.uniform((N, p_bin))
self._nodes.update({
"lin_pred_cts": self._lin_pred_cts,
"lin_pred_bin": self._lin_pred_bin
})
# factors
self._split = Split(parts={
"cts": self._lin_pred_cts,
"bin": self._lin_pred_bin
},
vector=self.parent)
self._model_cts = AddVariance(parent=self._lin_pred_cts,
child=self.child_cts,
variance=variance_cts)
if bin_model == "Logistic":
self._model_bin = Logistic(
parent=self._lin_pred_bin,
child=self.child_bin,
)
else:
self._model_bin = NoisyProbit(parent=self._lin_pred_bin,
child=self.child_bin,
variance=variance_bin)
self._factors.update({
"split": self._split,
"model_cts": self._model_cts,
"model_bin": self._model_bin
})
def __init__(self, vector: GaussianArray, parts: Dict[str, GaussianArray]):
super().__init__()
self._deterministic = True
self.vector = vector
shape_in = {n: p.shape() for n, p in parts.items()}
self.parts = parts
shape_out = vector.shape()
d = len(shape_out)
self.message_to_parts = {k: GaussianArray.uniform(s) for k, s in shape_in.items()}
self.message_to_vector = GaussianArray.uniform(shape_out)
size = [s[-1] for k, s in shape_in.items()]
begin = [0, *np.cumsum(size[:-1])]
self._size = [tuple([*shape_out[:-1], s]) for s in size]
self._begin = [tuple([*[0] * (d - 1), s]) for s in begin]
self._name = [k for k, s in shape_in.items()]
def to_v(self, x):
for k in x.keys():
x[k] = x[k] / self.message_to_x[k]
p = tf.concat([xx.precision() for k, xx in x.items()], -1)
mtp = tf.concat([xx.mean_times_precision() for k, xx in x.items()], -1)
self.message_to_v = GaussianArray.from_array_natural(p, mtp)
return self.message_to_v
def to_parent(self):
x = self.child / self.message_to_child
m = x.mean()
v = x.variance_safe() + self.variance.value()
message_to_parent = GaussianArray.from_array(m, v)
self.parent.update(self.message_to_parent, message_to_parent)
self.message_to_parent = message_to_parent
def to_sum(self, x):
x = x / self.message_to_x
self.message_to_sum = GaussianArray.from_array(
tf.math.reduce_sum(x.mean(), -1),
tf.math.reduce_sum(x.log_var(), -1)
)
return self.message_to_sum
def to_child(self):
mean = self.parent / self.message_to_parent
m = mean.mean()
v = mean.log_var() + self.variance.value()
message_to_child = GaussianArray.from_array(m, v)
self.child.update(self.message_to_child, message_to_child)
self.message_to_child = message_to_child
def __init__(self, child: BernoulliArray, parent: GaussianArray):
super().__init__()
self._deterministic = True
self.shape = child.shape()
self.child = child
self.parent = parent
self.message_to_child = BernoulliArray.uniform(self.shape)
self.message_to_parent = GaussianArray.uniform(self.shape)
def to_child(self):
x = self.parent / self.message_to_parent
message_to_child = GaussianArray.from_array(
tf.math.reduce_sum(x.mean(), -1),
tf.math.reduce_sum(x.log_var(), -1)
)
self.child.update(self.message_to_child, message_to_child)
self.message_to_child = message_to_child
def to_x(self, v):
v = v / self.message_to_v
p, mtp = v.natural()
for name, begin, size in zip(self._name, self._begin, self._size):
self.message_to_x[name] = GaussianArray.from_array_natural(
tf.slice(p, begin, size),
tf.slice(mtp, begin, size)
)
return self.message_to_x
def to_child(self):
x = self.parent / self.message_to_parent
weight = self.weight.value()
bias = self.bias.value()
m = tf.tensordot(x.mean(), weight, 1) + bias
v = tf.tensordot(x.log_var(), weight ** 2, 1)
message_to_child = GaussianArray.from_array(m, v)
self.child.update(self.message_to_child, message_to_child)
self.message_to_child = message_to_child
def to_vector(self):
parts = dict()
for name in self._name:
parts[name] = self.parts[name] / self.message_to_parts[name]
p = tf.concat([part.precision() for k, part in parts.items()], -1)
mtp = tf.concat([part.mean_times_precision() for k, part in parts.items()], -1)
message_to_vector = GaussianArray.from_array_natural(p, mtp)
self.vector.update(self.message_to_vector, message_to_vector)
self.message_to_vector = message_to_vector
def to_x(self, x, A):
# compute incoming message
from_x = x / self.message_to_x
# add noise
from_x = GaussianArray.from_array(
mean=from_x.mean(),
variance=from_x.log_var() + self.variance.value()
)
# compute outgoing message
A = A.proba() # TODO this assumes 0/1, need to fix later
stnr = from_x.mean() * tf.math.sqrt(from_x.precision()) * tf.cast(2*A - 1, tf.float32)
vf = tfp.distributions.Normal(0., 1.).prob(stnr) / tfp.distributions.Normal(0., 1.).cdf(stnr)
wf = vf * (stnr + vf)
m = from_x.mean() + tf.math.sqrt(from_x.log_var()) * vf * tf.cast(2 * A - 1, tf.float32)
v = from_x.log_var() * (1. - wf)
# add noise
v += self.variance.value()
self.message_to_x = GaussianArray.from_array(m, v)
return self.message_to_x
def to_product(self, x):
x = x / self.message_to_x
m0 = tf.expand_dims(x.mean(), 0)
m1 = tf.expand_dims(x.mean(), 1)
v0 = tf.expand_dims(x.log_var(), 0)
v1 = tf.expand_dims(x.log_var(), 1)
m = m0 * m1
v = m0 ** 2 * v1 + m1 ** 2 * v0 + v0 * v1
self.message_to_product = GaussianArray.from_array(m, v)
return self.message_to_product
def to_parts(self):
v = self.vector / self.message_to_vector
p, mtp = v.natural()
for name, begin, size in zip(self._name, self._begin, self._size):
message_to_part = GaussianArray.from_array_natural(
tf.slice(p, begin, size),
tf.slice(mtp, begin, size)
)
self.parts[name].update(self.message_to_parts[name], message_to_part)
self.message_to_parts[name] = message_to_part
def to_x(self, x, A):
# TODO: assumes 0/1 only, need to fix later with missing values
x = x / self.message_to_x
A = A.proba()
stnr = x.mean() * tf.math.sqrt(x.precision()) * tf.cast(2*A - 1, tf.float32)
vf = tfp.distributions.Normal(0., 1.).prob(stnr) / tfp.distributions.Normal(0., 1.).cdf(stnr)
wf = vf * (stnr + vf)
m = x.mean() + tf.math.sqrt(x.log_var()) * vf * tf.cast(2 * A - 1, tf.float32)
v = x.log_var() * (1. - wf)
self.message_to_x = GaussianArray.from_array(m, v)
return self.message_to_x
def to_x(self, x, sum):
sum = sum / self.message_to_sum
x = x / self.message_to_x
m = tf.expand_dims(sum.mean(), -1) - tf.math.reduce_sum(x.mean(), -1, keepdims=True) + x.mean()
v = tf.where(
x.is_uniform(),
np.inf,
tf.expand_dims(sum.log_var(), -1) + tf.math.reduce_sum(x.log_var(), -1, keepdims=True) - x.log_var()
)
self.message_to_x = GaussianArray.from_array(m, v)
return self.message_to_x
def to_result(self, x):
# compute incoming message
from_x = x / self.message_to_x
# add noise
from_x = GaussianArray.from_array(
mean=from_x.mean(),
variance=from_x.log_var() + self.variance.value()
)
# compute probability
proba = 1. - tfp.distributions.Normal(*from_x.mean_and_variance()).cdf(0.0)
return BernoulliArray.from_array(proba)
def to_parent(self):
child = self.child / self.message_to_child
p, mtp = child.natural()
p0 = tf.slice(p, [0, 0, 0], [-1, -1, 1])
mtp0 = tf.slice(mtp, [0, 0, 0], [-1, -1, 1])
p1 = tf.slice(p, [0, 0, 1], [-1, -1, 1])
mtp1 = tf.slice(mtp, [0, 0, 1], [-1, -1, 1])
p = tf.reduce_sum(p0, 0) + tf.reduce_sum(p1, 1)
mtp = tf.reduce_sum(mtp0, 0) + tf.reduce_sum(mtp1, 1)
message_to_parent = GaussianArray(p, mtp)
self.parent.update(self.message_to_parent, message_to_parent)
self.message_to_parent = message_to_parent
def to_child(self):
parent = self.parent / self.message_to_parent
p, mtp = parent.natural()
p0 = tf.tile(tf.expand_dims(p, 0), [self.N, 1, 1])
mtp0 = tf.tile(tf.expand_dims(mtp, 0), [self.N, 1, 1])
p1 = tf.tile(tf.expand_dims(p, 1), [1, self.N, 1])
mtp1 = tf.tile(tf.expand_dims(mtp, 1), [1, self.N, 1])
p = tf.concat([p0, p1], 2)
mtp = tf.concat([mtp0, mtp1], 2)
message_to_child = GaussianArray(p, mtp)
self.child.update(self.message_to_child, message_to_child)
self.message_to_child = message_to_child
def to_parent(self):
s = self.child / self.message_to_child
x = self.parent / self.message_to_parent
m = tf.expand_dims(s.mean(), -1) - tf.math.reduce_sum(x.mean(), -1, keepdims=True) + x.mean()
v = tf.where(
x.is_uniform(),
np.inf,
tf.expand_dims(s.log_var(), -1) + tf.math.reduce_sum(x.log_var(), -1, keepdims=True) - x.log_var()
)
message_to_parent = GaussianArray.from_array(m, v)
self.parent.update(self.message_to_parent, message_to_parent)
self.message_to_parent = message_to_parent
def to_elbo(self, x, A):
mean = x
x = GaussianArray.from_array(
mean=x.mean(),
variance=x.log_var() + self.variance.value()
)
elbo = mean.log_var() + mean.mean() ** 2
elbo += x.log_var() + x.mean() ** 2
elbo += -2. * x.mean() * mean.mean()
elbo /= self.variance.value()
# elbo += tf.math.log(2 * np.pi)
elbo += tf.math.log(self.variance.value())
return -.5 * tf.reduce_sum(elbo)
def to_x(self, x, result, weight, bias):
result = result / self.message_to_result
x = x / self.message_to_x
m = (tf.expand_dims(result.mean() - bias - tf.tensordot(x.mean(), weight, 1), 1) +
tf.expand_dims(x.mean(), -1) * tf.expand_dims(weight, 0)) / tf.expand_dims(weight, 0)
v = (tf.expand_dims(result.log_var() + tf.tensordot(x.log_var(), weight ** 2, 1), 1) -
tf.expand_dims(x.log_var(), -1) * tf.expand_dims(weight ** 2, 0)) / tf.expand_dims(weight ** 2, 0)
p = 1.0 / v
mtp = m * p
self.message_to_x = GaussianArray.from_array_natural(
tf.reduce_sum(p, -1),
tf.reduce_sum(mtp, -1)
)
return self.message_to_x
def to_parent(self):
x = self.parent / self.message_to_parent
A = self.child.proba()
A_safe = tf.where(tf.math.is_nan(A), 0.5, A)
stnr = x.mean() * tf.math.sqrt(x.precision()) * tf.cast(2 * A_safe - 1, tf.float32)
vf = tfp.distributions.Normal(0., 1.).prob(stnr) / tfp.distributions.Normal(0., 1.).cdf(stnr)
wf = vf * (stnr + vf)
m = x.mean() + tf.math.sqrt(x.log_var()) * vf * tf.cast(2 * A_safe - 1, tf.float32)
m = tf.where(A_safe == 0.5, 0., m)
v = x.log_var() * (1. - wf)
v = tf.where(A_safe == 0.5, 1.0e10, v)
message_to_parent = GaussianArray.from_array(m, v)
self.parent.update(self.message_to_parent, message_to_parent)
self.message_to_parent = message_to_parent