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