def precompute(self): #definition of q(u) M = self.Z.shape[0] self.mu = tf.Variable(0.3 * tf.random.normal([M, 1])) self.scale = tf.Variable(np.tril(0.1 * np.random.randn(M, M) + 1.0 * np.eye(M)), dtype=tf.float32) #self.mu = tf.Variable(0*tf.random.normal([M,1])) #self.scale = tf.Variable(5*np.eye(M),dtype=tf.float32) #tf.Variable(0.2*np.tril(1.0*np.random.randn(M,M)+1.0*np.eye(M)),dtype=tf.float32) if self.likemodel == 'distribution' or self.likemodel == 'process': Mlike = self.Zlike.shape[0] self.mulike = tf.Variable(0.0001 * tf.random.normal([Mlike, 1])) mu_u = tf.Variable(np.full([Mlike], -12), dtype=tf.float32) cov_u = tf.Variable(self.klike.matrix(self.Zlike, self.Zlike), dtype=tf.float32) self.pulike = tfd.MultivariateNormalFullCovariance( mu_u, cov_u + np.eye(cov_u.shape[0]) * self.jitter) else: self.mulike = None if self.likemodel == 'process': self.scalelike = tf.Variable(1e-10 * np.eye(Mlike), dtype=tf.float32) else: self.scalelike = None #parameters for p(u) mu_u = tf.zeros([M], dtype=tf.float32) cov_u = tf.Variable(self.k.matrix(self.Z, self.Z), dtype=tf.float32) self.pu = tfd.MultivariateNormalFullCovariance( mu_u, cov_u + np.eye(cov_u.shape[0]) * self.jitter)
def demo_gmm_log_pdf(z): norm1 = tfd.MultivariateNormalFullCovariance( loc=demo_gmm_mu1, covariance_matrix=demo_gmm_sigma1) norm2 = tfd.MultivariateNormalFullCovariance( loc=demo_gmm_mu2, covariance_matrix=demo_gmm_sigma2) return math.log(0.5 * norm1.prob(z) + 0.5 * norm2.prob(z))
def conditional_distribution(self, x, slice_in, slice_out): marginal_in = ds.MixtureSameFamily( mixture_distribution=ds.Categorical(probs=self._priors), components_distribution=ds.MultivariateNormalFullCovariance( loc=self._locs[:, slice_in], covariance_matrix=self._covs[:, slice_in, slice_in])) p_k_in = ds.Categorical(logits=tf_utils.log_normalize( marginal_in.components_distribution.log_prob(x[:, None]) + marginal_in.mixture_distribution.logits[None], axis=1)) sigma_in_out = self._covs[:, slice_in, slice_out] inv_sigma_in_in = tf.linalg.inv(self._covs[:, slice_in, slice_in]) inv_sigma_out_in = tf.matmul(sigma_in_out, inv_sigma_in_in, transpose_a=True) A = inv_sigma_out_in b = self._locs[:, slice_out] - tf.matmul( inv_sigma_out_in, self._locs[:, slice_in, None])[:, :, 0] cov_est = (self._covs[:, slice_out, slice_out] - tf.matmul(inv_sigma_out_in, sigma_in_out)) ys = tf.einsum('aij,bj->abi', A, x) + b[:, None] p_out_in_k = ds.MultivariateNormalFullCovariance( tf.transpose(ys, perm=(1, 0, 2)), cov_est) return ds.MixtureSameFamily(mixture_distribution=p_k_in, components_distribution=p_out_in_k)
def get_samples(self, mu, scale, num=100): """ Get samples of the function components for every observation pair in X. Returns a num x N x (C*2) matrix, where num = number of samples N = number of observation pairs C = number of components So for the tensor that is returned, the last dimension consists of the pairs of sensors, with each pair being one of the C components. If scale is set to None, then we return a single sample, of the posterior mean (i.e. we assume a dirac q(f). Returns 1 x N x (C*2) matrix. """ qf_mu, qf_cov = self.get_qf(mu, scale) if scale is None: return tf.transpose(tf.reshape( qf_mu, [2 * self.C, self.Npairs]))[None, :, :] batched_mu = tf.transpose(tf.reshape(qf_mu, [2 * self.C, self.Npairs])) batched_cov = [] for ni in range(0, self.Npairs * self.C * 2, self.Npairs): innerbcov = [] for nj in range(0, self.Npairs * self.C * 2, self.Npairs): innerbcov.append( tf.linalg.diag_part(qf_cov[ni:(ni + self.Npairs), nj:(nj + self.Npairs)])) batched_cov.append(innerbcov) samps = tfd.MultivariateNormalFullCovariance( batched_mu, tf.transpose(batched_cov) + tf.eye(2 * self.C) * self.jitter).sample(num) return samps
def get_samples_one_sensor(self, mu, scale, num=100): """ Get samples of the function components for a sensor. Returns a num x N x (C) matrix, where num = number of samples N = number of observation pairs C = number of components So for the tensor that is returned, the last dimension consists of the pairs of sensors, with each pair being one of the C components. """ qf_mu, qf_cov = self.get_qf(mu, scale) if scale is None: return tf.transpose(tf.reshape(qf_mu, [self.C, self.N]))[None, :, :] batched_mu = tf.transpose(tf.reshape(qf_mu, [self.C, self.N])) batched_cov = [] for ni in range(0, self.N * self.C, self.N): innerbcov = [] for nj in range(0, self.N * self.C, self.N): innerbcov.append( tf.linalg.diag_part(qf_cov[ni:(ni + self.N * 2), nj:(nj + self.N)])) batched_cov.append(innerbcov) samps = tfd.MultivariateNormalFullCovariance( batched_mu, tf.transpose(batched_cov) + tf.eye(self.C) * self.jitter).sample(num) return samps
def test_factored_joint_mvn_diag_full(self): batch_shape = [3, 2] mvn1 = tfd.MultivariateNormalDiag(loc=tf.zeros(batch_shape + [3]), scale_diag=tf.ones(batch_shape + [3])) mvn2 = tfd.MultivariateNormalFullCovariance( loc=tf.ones(batch_shape + [2]), covariance_matrix=(tf.ones(batch_shape + [2, 2]) * [[5., -2], [-2, 3.1]])) joint = sts_util.factored_joint_mvn([mvn1, mvn2]) self.assertEqual( self.evaluate(joint.event_shape_tensor()), self.evaluate(mvn1.event_shape_tensor() + mvn2.event_shape_tensor())) joint_mean_ = self.evaluate(joint.mean()) self.assertAllEqual(joint_mean_[..., :3], self.evaluate(mvn1.mean())) self.assertAllEqual(joint_mean_[..., 3:], self.evaluate(mvn2.mean())) joint_cov_ = self.evaluate(joint.covariance()) self.assertAllEqual(joint_cov_[..., :3, :3], self.evaluate(mvn1.covariance())) self.assertAllEqual(joint_cov_[..., 3:, 3:], self.evaluate(mvn2.covariance()))
def _get_prior(self): if self.prior is None: # Compute kernel matrices for each latent dimension kernel_matrices = [] for i in range(self.kernel_scales): if self.kernel == "rbf": kernel_matrices.append(rbf_kernel(self.time_length, self.length_scale / 2**i)) elif self.kernel == "diffusion": kernel_matrices.append(diffusion_kernel(self.time_length, self.length_scale / 2**i)) elif self.kernel == "matern": kernel_matrices.append(matern_kernel(self.time_length, self.length_scale / 2**i)) elif self.kernel == "cauchy": kernel_matrices.append(cauchy_kernel(self.time_length, self.sigma, self.length_scale / 2**i)) # Combine kernel matrices for each latent dimension tiled_matrices = [] total = 0 for i in range(self.kernel_scales): if i == self.kernel_scales-1: multiplier = self.latent_dim - total else: multiplier = int(np.ceil(self.latent_dim / self.kernel_scales)) total += multiplier tiled_matrices.append(tf.tile(tf.expand_dims(kernel_matrices[i], 0), [multiplier, 1, 1])) kernel_matrix_tiled = np.concatenate(tiled_matrices) assert len(kernel_matrix_tiled) == self.latent_dim self.prior = tfd.MultivariateNormalFullCovariance( loc=tf.zeros([self.latent_dim, self.time_length], dtype=tf.float32), covariance_matrix=kernel_matrix_tiled) return self.prior
def get_mvn_from_transformation(self, Input): with tf.variable_scope(self.name, reuse=tf.AUTO_REUSE): sigma = None mu = self.transformation.transform(Input) if isinstance(mu, tuple): assert len( mu) == 2, "output of {} should contain 2 elements".format( self.transformation.name) mu, sigma = mu sigma_con = self.get_sigma(mu) if sigma is None: mvn = tfd.MultivariateNormalDiag(mu, sigma_con, validate_args=True, allow_nan_stats=False) else: if len(sigma.shape.as_list()) == len(mu.shape.as_list()): sigma = sigma_con + 0.1 * sigma mvn = tfd.MultivariateNormalDiag(mu, sigma, validate_args=True, allow_nan_stats=False) else: sigma = tf.diag(sigma_con) + 0.1 * sigma mvn = tfd.MultivariateNormalFullCovariance( mu, sigma, validate_args=True, allow_nan_stats=False) return mvn
def output_function(self, state): params = dense_layer(state.h3, self.output_units, scope='gmm', reuse=tf.compat.v1.AUTO_REUSE) pis, mus, sigmas, rhos, es = self._parse_parameters(params) mu1, mu2 = tf.split(mus, 2, axis=1) mus = tf.stack([mu1, mu2], axis=2) sigma1, sigma2 = tf.split(sigmas, 2, axis=1) covar_matrix = [ tf.square(sigma1), rhos * sigma1 * sigma2, rhos * sigma1 * sigma2, tf.square(sigma2) ] covar_matrix = tf.stack(covar_matrix, axis=2) covar_matrix = tf.reshape( covar_matrix, (self.batch_size, self.num_output_mixture_components, 2, 2)) mvn = tfd.MultivariateNormalFullCovariance( loc=mus, covariance_matrix=covar_matrix) b = tfd.Bernoulli(probs=es) c = tfd.Categorical(probs=pis) sampled_e = b.sample() sampled_coords = mvn.sample() sampled_idx = c.sample() idx = tf.stack([tf.range(self.batch_size), sampled_idx], axis=1) coords = tf.gather_nd(sampled_coords, idx) return tf.concat([coords, tf.cast(sampled_e, tf.float32)], axis=1)
def get_q(self, h_stack, y_t_ft): """ calculate x_t ~ q(*|h_t-1, y_t_ft) h_stack.shape = (n_particles, batch_size, Dh) y_t_ft.shape = (batch_size, Dy_1) """ with tf.variable_scope(self.variable_scope + '/get_q'): y_t_ft_expanded = tf.expand_dims(y_t_ft, axis = 0, name = 'y_t_ft_expanded') # y_t_ft_expanded.shape = (1, batch_size, Dy_1) y_t_ft_tiled = tf.tile(y_t_ft_expanded, (self.n_particles, 1, 1), name = 'y_t_ft_tiled') # y_t_ft_tiled.shape = (n_paticles, batch_size, Dy_1) h_y_concat = tf.concat((h_stack, y_t_ft_tiled), axis = 2, name = 'h_y_concat') # h_y_concat.shape = (n_paticles, batch_size, Dh + Dy_1) mu = fully_connected(h_y_concat, self.Dx, weights_initializer=xavier_initializer(uniform=False), activation_fn = None, reuse = tf.AUTO_REUSE, scope = "mu") # mu.shape = (n_paticles, batch_size, Dx) sigma = fully_connected(h_y_concat, self.Dx, weights_initializer=xavier_initializer(uniform=False), biases_initializer=tf.constant_initializer(0.6), activation_fn = tf.nn.softplus, reuse = tf.AUTO_REUSE, scope = "sigma") + self.sigma_cons # sigma.shape = (n_paticles, batch_size, Dx) q = tfd.MultivariateNormalFullCovariance(loc = mu, covariance_matrix = tf.matrix_diag(sigma), name = "q") return q
def __init__(self, priors, locs, covs, init_params='kmeans', warm_start=True, reg_cov=1e-6, bayesian=False): self._priors = priors self._locs = locs self._covs = covs self._covs_par = covs if tf.__version__[0] == '1': self._k = priors.shape[0].value self._n = locs.shape[-1].value else: self._k = priors.shape[0] self._n = locs.shape[-1] try: m = BayesianGaussianMixture if bayesian else GaussianMixture self._sk_gmm = m(self._k, 'full', n_init=1, init_params=init_params, warm_start=warm_start, reg_covar=reg_cov) except: print( "Sklearn not install, cannot perform MLE of Gaussian mixture model" ) self._priors_ml = tf.compat.v1.placeholder(tf.float32, (self._k, )) self._locs_ml = tf.compat.v1.placeholder(tf.float32, (self._k, self._n)) self._covs_ml = tf.compat.v1.placeholder(tf.float32, (self._k, self._n, self._n)) self._ml_assign_op = [ self._priors.assign(self._priors_ml), self._locs.assign(self._locs_ml), self._covs.assign(self._covs_ml) ] ds.MixtureSameFamily.__init__( self, mixture_distribution=ds.Categorical(probs=self._priors), components_distribution=ds.MultivariateNormalFullCovariance( loc=self._locs, covariance_matrix=self._covs)) self._init = False
def _create_prior(self): if self._learn_prior == LEARN_PRIOR_MAF: self._prior = IndependentChannelsTransformedDistribution( [1, self.time_dim_latent], self._flow_params, self.time_dim_latent, self._latent_channels ) elif self._variational_layer == GAUSSIAN_DIAGONAL: dim = tf.shape(self._gaussian_model_latent.mean())[1:3] zeros = tf.zeros(shape=dim) ones = self._alpha * tf.ones(shape=dim) self._prior = tfd.MultivariateNormalDiag(loc=zeros, scale_diag=ones, name="prior") elif self._variational_layer == GAUSSIAN_TRI_DIAGONAL: dim = tf.shape(self._gaussian_model_latent.mean())[1:3] ch_z = self._gaussian_model_latent.mean().shape[1] zeros = tf.zeros(shape=dim) ones = tf.stack([tf.eye(dim[1]) for i in range(ch_z)]) twos = self._alpha * 0.2 * tf.stack([tf.eye(dim[1] - 1) for i in range(ch_z)]) twos_big = tf.pad(twos, [[0, 0], [1, 0], [0, 1]], mode='CONSTANT') cov = ones + twos_big + tf.transpose(twos_big, perm=[0, 2, 1]) self._prior = tfd.MultivariateNormalFullCovariance(loc=zeros, covariance_matrix=cov, name="prior") elif self._variational_layer == GAUSSIAN_TRI_DIAGONAL_PRECISION: dim = tf.shape(self._gaussian_model_latent.mean())[1:3] ch_z = self._gaussian_model_latent.mean().shape[1] zeros = tf.zeros(shape=dim) ones = tf.stack([tf.eye(dim[1]) for i in range(ch_z)]) twos = self._alpha * tf.stack([tf.eye(dim[1] - 1) for i in range(ch_z)]) twos_big = tf.pad(twos, [[0, 0], [1, 0], [0, 1]], mode='CONSTANT') L = ones + twos_big L_T = tf.transpose(L, perm=[0, 2, 1]) prec = tf.matmul(L, L_T) cov = tf.linalg.inv(prec) self._prior = tfd.MultivariateNormalFullCovariance(loc=zeros, covariance_matrix=cov, name="prior") else: raise ValueError("specified string is not a suitable variational layer: %s" % str(self._variational_layer))
def concs_r_rv(self): """ Random variable describing the receptor concentrations of the cell.""" if (self.analytical == False) and (self.sampled == False): self.sample() if (self.analytical == True) and (self.finished_select_non_zero == False): self.select_non_zero() return copy.deepcopy( tfd.MultivariateNormalFullCovariance( loc=np.array(self.concs[2, :], dtype=np.float32), covariance_matrix=np.array(self.concs_cov[2, :, :], dtype=np.float32)))
def precompute(self): #definition of q(u) M = self.Z.shape[0] self.mu = tf.Variable(0.3 * tf.random.normal([M, 1])) self.scale = tf.Variable( 0.05 * np.tril(1.0 * np.random.randn(M, M) + 1.0 * np.eye(M)), dtype=tf.float32) if self.likemodel == 'distribution' or self.likemodel == 'process': Mlike = self.Zlike.shape[0] self.mulike = tf.Variable(0.0001 * tf.random.normal([Mlike, 1])) mu_u = tf.Variable(np.full([Mlike], -12), dtype=tf.float32) cov_u = tf.Variable(self.klike.matrix(self.Zlike, self.Zlike), dtype=tf.float32) self.pulike = tfd.MultivariateNormalFullCovariance( mu_u, cov_u + np.eye(cov_u.shape[0]) * self.jitter) #0.8 self.smlike = SparseModelNoMiniBatch(self.Xlike, self.Zlike, 1, self.k) else: self.mulike = None if self.likemodel == 'process': self.scalelike = tf.Variable(1e-10 * np.eye(Mlike), dtype=tf.float32) else: self.scalelike = None #parameters for p(u) mu_u = tf.zeros([M], dtype=tf.float32) cov_u = tf.Variable(self.k.matrix(self.Z, self.Z), dtype=tf.float32) self.pu = tfd.MultivariateNormalFullCovariance( mu_u, cov_u + np.eye(cov_u.shape[0]) * self.jitter) self.ref = tf.gather( self.refsensor, tf.transpose( tf.reshape(self.X[:(2 * self.N), 1:2].astype(int), [2, self.N]))) self.sm = SparseModelNoMiniBatch(self.X, self.Z, self.C, self.k)
def sample_from_2_dist(self, dist1, dist2, d1_input, d2_input, sample_size=()): d1_mvn = dist1.get_mvn(d1_input) d2_mvn = dist2.get_mvn(d2_input) if isinstance(d1_mvn, tfd.MultivariateNormalDiag) and isinstance(d2_mvn, tfd.MultivariateNormalDiag): d1_mvn_mean, d1_mvn_cov = d1_mvn.mean(), d1_mvn.stddev() d2_mvn_mean, d2_mvn_cov = d2_mvn.mean(), d2_mvn.stddev() d1_mvn_cov_inv, d2_mvn_cov_inv = 1 / d1_mvn_cov, 1 / d2_mvn_cov combined_cov = 1 / (d1_mvn_cov_inv + d2_mvn_cov_inv) combined_mean = combined_cov * (d1_mvn_cov_inv * d1_mvn_mean + d2_mvn_cov_inv * d2_mvn_mean) mvn = tfd.MultivariateNormalDiag(combined_mean, combined_cov, validate_args=True, allow_nan_stats=False) else: if isinstance(d1_mvn, tfd.MultivariateNormalDiag): d1_mvn_mean, d1_mvn_cov = d1_mvn.mean(), tf.diag(d1_mvn.stddev()) elif isinstance(d1_mvn, tfd.MultivariateNormalFullCovariance): d1_mvn_mean, d1_mvn_cov = d1_mvn.mean(), d1_mvn.covariance() if isinstance(d2_mvn, tfd.MultivariateNormalDiag): d2_mvn_mean, d2_mvn_cov = d2_mvn.mean(), tf.diag(d2_mvn.stddev()) elif isinstance(d2_mvn, tfd.MultivariateNormalFullCovariance): d2_mvn_mean, d2_mvn_cov = d2_mvn.mean(), d2_mvn.covariance() if len(d1_mvn_cov.shape.as_list()) == 2: d1_mvn_cov = tf.expand_dims(d1_mvn_cov, axis=0) d1_mvn_cov_inv, d2_mvn_cov_inv = tf.linalg.inv(d1_mvn_cov), tf.linalg.inv(d2_mvn_cov) combined_cov = tf.linalg.inv(d1_mvn_cov_inv + d2_mvn_cov_inv) perm = list(range(len(combined_cov.shape))) perm[-2], perm[-1] = perm[-1], perm[-2] combined_cov = (combined_cov + tf.transpose(combined_cov, perm=perm)) / 2 combined_mean = tf.matmul(combined_cov, tf.matmul(d1_mvn_cov_inv, tf.expand_dims(d1_mvn_mean, axis=-1)) + tf.matmul(d2_mvn_cov_inv, tf.expand_dims(d2_mvn_mean, axis=-1)) ) combined_mean = tf.squeeze(combined_mean, axis=-1) mvn = tfd.MultivariateNormalFullCovariance(combined_mean, combined_cov, validate_args=True, allow_nan_stats=False) X = mvn.sample(sample_size) q_t_log_prob = mvn.log_prob(X) f_t_log_prob = d1_mvn.log_prob(X) return X, q_t_log_prob, f_t_log_prob
def __init__(self, loc=None, covariance_matrix=None, scale_obs=0.1, psi=None, fast_sample=True, validate_args=False, allow_nan_stats=True, name="ProMP"): self._psi = psi self._psi_op = tf.linalg.LinearOperatorFullMatrix(psi) self._loc_w = loc self._cov_w = covariance_matrix self._mvn_w = ds.MultivariateNormalFullCovariance( loc=self._loc_w, covariance_matrix=self._cov_w) # (psi T cov psi ) T = (psi T cov )T # _loc = tf.linalg.matvec(psi, loc, transpose_a=True) # _cov = tf.linalg.matmul(tf.linalg.matmul(psi, covariance_matrix, transpose_a=True), # psi) + \ # tf.eye(psi.shape[1].value) * scale_obs ** 2 _loc = self._psi_op.matvec(loc, adjoint=True) _cov = self._psi_op.matmul(self._psi_op.matmul(covariance_matrix, adjoint=True), adjoint=True, adjoint_arg=True) self._mvn_obs = ds.MultivariateNormalDiag( loc=tf.zeros_like(_loc), scale_diag=scale_obs * tf.ones(psi.shape[1].value), ) self._fast_sample = fast_sample if _cov.shape.ndims == 2: _cov += tf.eye(psi.shape[1].value) * scale_obs**2 elif _cov.shape.ndims == 3: _cov += tf.eye(psi.shape[1].value)[None] * scale_obs**2 else: raise NotImplementedError ds.MultivariateNormalFullCovariance.__init__(self, loc=_loc, covariance_matrix=_cov)
def _build(self, inputs): # TODO this could be adapter to be compatible with non linear layers #mean, covariance, scale = self.create_mean_n_cov_layers(inputs) inputs = tf.layers.flatten(inputs) linear_mean = snt.Linear(output_size=output_size) mean = linear_mean(inputs) pdb.set_trace() linear_cov = snt.Linear(output_size=[output_size, output_size]) cov = linear_mean(inputs) # TODO this neds to be adapter to the new covariance #self.set_contractive_regularizer(mean, covariance, # self._contractive_regularizer_inputs, # self._contractive_regularizer_tuple, # self._contractive_collection_network_str) output_distribution = tfd.MultivariateNormalFullCovariance(loc=mean, covariance_matrix=cov)
def log_prob(self, us, xis, i=None): """ :param xis: [horizon + 1, batch_size, xi_dim] :param us: [horizon, batch_size, xi_dim] :param i: :return: """ if i is not None: raise NotImplementedError # compute u for each timestep and batch b, as size [horzon, batch_size, u_dim] u_locs = -tf.einsum('aij,abj->abi', self._Ks[:self.horizon], xis[:self.horizon]) + \ tf.einsum('aij,aj->ai', self._Kvs[:self.horizon], self._vs[1:self.horizon+1])[:, None] u_covs = self._Qs[:self.horizon][:, None] return ds.MultivariateNormalFullCovariance(loc=u_locs, covariance_matrix=u_covs).log_prob(us)
def reduce_mvn_ds(mvns): """ Perform moment matching mvns : list of mvn :return: """ # make h [..., 1], multiply and reduce locs = tf.stack([mvn.loc for mvn in mvns]) loc = tf.reduce_mean(locs, axis=0) dlocs = locs - tf.expand_dims(loc, axis=0) cov_locs = tf.matmul(tf.expand_dims(dlocs, axis=-1), tf.expand_dims(dlocs, axis=-2)) covs = tf.stack([mvn.covariance() for mvn in mvns]) cov = tf.reduce_mean(covs + cov_locs, axis=0) return distributions.MultivariateNormalFullCovariance(loc, cov)
def get_f(self, h_stack): """ calculate x_t ~ f(*|h_t-1) h_stack.shape = (n_particles, batch_size, Dh) """ with tf.variable_scope(self.variable_scope + '/get_f'): mu = fully_connected(h_stack, self.Dx, weights_initializer=xavier_initializer(uniform=False), activation_fn = None, reuse = tf.AUTO_REUSE, scope = "mu") # mu.shape = (n_paticles, batch_size, Dx) sigma = fully_connected(h_stack, self.Dx, weights_initializer=xavier_initializer(uniform=False), biases_initializer=tf.constant_initializer(0.6), activation_fn = tf.nn.softplus, reuse = tf.AUTO_REUSE, scope = "sigma") + self.sigma_cons # sigma.shape = (n_paticles, batch_size, Dx) f = tfd.MultivariateNormalFullCovariance(loc = mu, covariance_matrix = tf.matrix_diag(sigma), name = "f") return f
def __init__(self, mixture, marginal_slice=None, bayesian=False): self._logits = param.make_logits_from_value(mixture.weights_) if bayesian: mixture = filter_unused(mixture) if marginal_slice is not None: self._locs = param.make_loc_from_value( mixture.means_[:, marginal_slice]) self._covs = param.make_cov_from_value( mixture.covariances_[:, marginal_slice, marginal_slice]) else: self._locs = param.make_loc_from_value(mixture.means_) self._covs = param.make_cov_from_value(mixture.covariances_) self._priors = tf.math.exp(self._logits) ds.MixtureSameFamily.__init__( self, mixture_distribution=ds.Categorical(probs=self._priors), components_distribution=ds.MultivariateNormalFullCovariance( loc=self._locs, covariance_matrix=self._covs))
def get_g(self, h_stack, x_t): """ calculate y_t ~ g(*|h_t-1, x_t) h_stack.shape = (n_particles, batch_size, Dh) x_t.shape = (n_particles, batch_size, Dx) """ x_t_ft = self.get_x_ft(x_t) # x_t_ft.shape = (n_particles, batch_size, Dx_1) with tf.variable_scope(self.variable_scope + '/get_g'): h_x_concat = tf.concat((h_stack, x_t_ft), axis = 2, name = 'h_x_concat') mu = fully_connected(h_x_concat, self.Dy, weights_initializer=xavier_initializer(uniform=False), activation_fn = None, reuse = tf.AUTO_REUSE, scope = "mu") # mu.shape = (n_paticles, batch_size, Dx) sigma = fully_connected(h_x_concat, self.Dy, weights_initializer=xavier_initializer(uniform=False), biases_initializer=tf.constant_initializer(0.6), activation_fn = tf.nn.softplus, reuse = tf.AUTO_REUSE, scope = "sigma") + self.sigma_cons # sigma.shape = (n_paticles, batch_size, Dx) g = tfd.MultivariateNormalFullCovariance(loc = mu, covariance_matrix = tf.matrix_diag(sigma), name = "g") return g
def _init_distribution(conditions, **kwargs): loc, covariance_matrix = conditions["loc"], conditions["covariance_matrix"] return tfd.MultivariateNormalFullCovariance( loc=loc, covariance_matrix=covariance_matrix, **kwargs )
def tensorflow_model(self, pars): """Output tensorflow probability model object, to be combined together and sampled from. pars - dictionary of signal and nuisance parameters (tensors, constant or Variable) """ # Need to construct these shapes to match the event_shape, batch_shape, sample_shape # semantics of tensorflow_probability. cov_order = self.get_cov_order() small = 1e-10 #print("pars:",pars) tfds = {} # Determine which SRs participate in the covariance matrix if self.cov is not None: cov = tf.constant(self.cov, dtype=c.TFdtype) cov_diag = tf.constant( [self.cov[k][k] for k in range(len(self.cov))]) # Select which systematic to use, depending on whether SR participates in the covariance matrix bsys_tmp = [ np.sqrt(self.cov_diag[cov_order.index(sr)]) if self.in_cov[i] else self.SR_b_sys[i] for i, sr in enumerate(self.SR_names) ] else: bsys_tmp = self.SR_b_sys[:] # Prepare input parameters #print("input pars:",pars) s = pars[ 's'] * self.s_scaling # We "scan" normalised versions of s, to help optimizer theta = pars[ 'theta'] * self.theta_scaling # We "scan" normalised versions of theta, to help optimizer #print("de-scaled pars: s :",s) #print("de-scaled pars: theta:",theta) theta_safe = theta #print("theta_safe:", theta_safe) #print("rate:", s+b+theta_safe) # Expand dims of internal parameters to match input pars. Right-most dimension is the 'event' dimension, # i.e. parameters for each independent Poisson distribution. The rest go into batch_shape. if s.shape == (): n_batch_dims = 0 else: n_batch_dims = len(s.shape) - 1 new_dims = [1 for i in range(n_batch_dims)] b = tf.constant(self.SR_b, dtype=c.TFdtype) bsys = tf.constant(bsys_tmp, dtype=c.TFdtype) if n_batch_dims > 0: b = tf.reshape(b, new_dims + list(b.shape)) bsys = tf.reshape(bsys, new_dims + list(bsys.shape)) # Poisson model poises0 = tfd.Poisson( rate=tf.abs(s + b + theta_safe) + c.reallysmall ) # Abs works to constrain rate to be positive. Might be confusing to interpret BF parameters though. # Treat SR batch dims as event dims poises0i = tfd.Independent(distribution=poises0, reinterpreted_batch_ndims=1) tfds["n"] = poises0i # Multivariate background constraints if self.cov is not None: #print("theta_safe:",theta_safe) #print("covi:",self.covi) theta_cov = tf.gather(theta_safe, self.covi, axis=-1) #print("theta_cov:",theta_cov) cov_nuis = tfd.MultivariateNormalFullCovariance( loc=theta_cov, covariance_matrix=cov) tfds["x_cov"] = cov_nuis #print("str(cov_nuis):", str(cov_nuis)) # Remaining uncorrelated background constraints if np.sum(~self.in_cov) > 0: nuis0 = tfd.Normal(loc=theta_safe[..., ~self.in_cov], scale=bsys[..., ~self.in_cov]) # Treat SR batch dims as event dims nuis0i = tfd.Independent(distribution=nuis0, reinterpreted_batch_ndims=1) tfds["x_nocov"] = nuis0i else: # Only have uncorrelated background constraints nuis0 = tfd.Normal(loc=theta_safe, scale=bsys) # Treat SR batch dims as event dims nuis0i = tfd.Independent(distribution=nuis0, reinterpreted_batch_ndims=1) tfds["x"] = nuis0i #print("hello3") return tfds #, sample_layout, sample_count
def _init_distribution(conditions): mu, cov = conditions["mu"], conditions["cov"] return tfd.MultivariateNormalFullCovariance(loc=mu, covariance_matrix=cov)
def __init__(self, log_unnormalized_prob, gmm=None, k=10, loc=0., std=1., ndim=None, loc_tril=None, samples=20, temp=1., cov_type='diag', loc_scale=1., priors_scale=1e1): """ :param log_unnormalized_prob: Unnormalized log density to estimate :type log_unnormalized_prob: a tensorflow function that takes [batch_size, ndim] as input and returns [batch_size] :param gmm: :param k: number of components for GMM approximation :param loc: for initialization, mean :param std: for initialization, standard deviation :param ndim: """ self.log_prob = log_unnormalized_prob self.ndim = ndim self.temp = temp if gmm is None: assert ndim is not None, "If no gmm is defined, should give the shape of x" if cov_type == 'diag': _log_priors_var = tf.Variable(1. / priors_scale * log_normalize(tf.ones(k))) log_priors = priors_scale * _log_priors_var if isinstance(loc, tf.Tensor) and loc.shape.ndims == 2: _locs_var = tf.Variable(1. / loc_scale * loc) locs = loc_scale * _locs_var else: _locs_var = tf.Variable( 1. / loc_scale * tf.random.normal((k, ndim), loc, std)) locs = loc_scale * _locs_var log_std_diags = tf.Variable(tf.log(std/k * tf.ones((k, ndim)))) self._opt_params = [_log_priors_var, _locs_var, log_std_diags] gmm = _distributions.MixtureSameFamily( mixture_distribution=_distributions.Categorical(logits=log_priors), components_distribution=_distributions.MultivariateNormalDiag( loc=locs, scale_diag=tf.math.exp(log_std_diags) ) ) elif cov_type == 'full': _log_priors_var = tf.Variable(1./priors_scale * log_normalize(tf.ones(k))) log_priors = priors_scale * _log_priors_var if isinstance(loc, tf.Tensor) and loc.shape.ndims == 2: _locs_var = tf.Variable(1. / loc_scale * loc) locs = loc_scale * _locs_var else: _locs_var = tf.Variable(1./loc_scale * tf.random.normal((k, ndim), loc, std)) locs = loc_scale * _locs_var loc_tril = loc_tril if loc_tril is not None else std/k # tril_cov = tf.Variable(loc_tril ** 2 * tf.eye(ndim, batch_shape=(k, ))) tril_cov = tf.Variable(tf1.log(loc_tril) * tf.eye(ndim, batch_shape=(k, ))) covariance = tf.linalg.expm(tril_cov + tf1.matrix_transpose(tril_cov)) # self._opt_params = [_log_priors_var, _locs_var, tril_cov] gmm = _distributions.MixtureSameFamily( mixture_distribution=_distributions.Categorical(logits=log_priors), components_distribution=_distributions.MultivariateNormalFullCovariance( loc=locs, covariance_matrix=covariance ) ) else: raise ValueError("Unrecognized covariance type") self.k = k self.num_samples = samples self.gmm = gmm
def get_u_samples(self, xi, i=0): loc = -tf.einsum('ij,aj->ai', self._Ks[i], xi) + tf.einsum('ij,j->i', self._Kvs[i], self._vs[i+1])[None] cov = self._Qs[i] return ds.MultivariateNormalFullCovariance(loc=loc, covariance_matrix=cov).sample(1)[0]
def entropy(self): return tf.reduce_sum(ds.MultivariateNormalFullCovariance( covariance_matrix=self._Qs).entropy())
def build_graph(self, paradigm, weights, parameters, data=None, rho_init=.5, lambd=1.): self.graph = tf.Graph() with self.graph.as_default(): self.parameters_ = tf.constant(parameters) # n_timepoints x n_stim_dimensions x n_basis functions x n_voxels self.paradigm_ = tf.constant(paradigm.values[..., np.newaxis, np.newaxis], dtype=tf.float32, name='paradigm') self.weights_ = tf.Variable(weights.values[np.newaxis, np.newaxis, ...], dtype=tf.float32, name='basis_weights') self.build_basis_function() # n_timepoints x n_voxels self.predictions_ = tf.squeeze( tf.tensordot(self.basis_predictions_, self.weights_, (1, 2))) # Simulation self.noise_ = tf.placeholder(tf.float32, shape=(1, None), name='noise') n_timepoints, n_voxels = tf.shape(self.paradigm_)[0], tf.shape( self.weights_)[-1] noise = tf.random_normal(shape=(n_timepoints, n_voxels), mean=0.0, stddev=self.noise_, dtype=tf.float32) self.noisy_predictions_ = self.predictions_ + noise # Data and residuals if data is not None: data = pd.DataFrame(data) self.data_ = tf.constant(data.values, name='data') self.residuals_ = self.data_ - self.predictions_ # Residual model self.rho_trans = tf.Variable(rho_init, dtype=tf.float32, name='rho_trans') self.rho_ = tf.math.sigmoid(self.rho_trans, name='rho') self.tau_trans = tf.Variable(_inverse_softplus( data.std().values[:, np.newaxis]), name='tau_trans') self.tau_ = _softplus_tensor(self.tau_trans, name='tau') self.sigma2_trans = tf.Variable(0., dtype=tf.float32, name='sigma2_trans') self.sigma2_ = _softplus_tensor(self.sigma2_trans, name='sigma2') sigma0 = self.rho_ * tf.tensordot(self.tau_, tf.transpose(self.tau_), axes=1) + \ (1 - self.rho_) * tf.linalg.tensor_diag(tf.squeeze(self.tau_**2)) + \ self.sigma2_ * tf.squeeze(tf.tensordot(self.weights_, self.weights_, axes=(-2, -2))) self.empirical_covariance_matrix_ = tf.constant( data.cov().values.astype(np.float32), name='empirical_covariance_matrix') self.sigma_ = lambd * sigma0 + \ (1 - lambd) * self.empirical_covariance_matrix_ self.residual_dist = tfd.MultivariateNormalFullCovariance( tf.zeros(data.shape[1]), self.sigma_) self.likelihood_ = self.residual_dist.log_prob(self.residuals_)
def __call__(self, prediction_horizon, view = False): # self.reset_variables() rewards = 0 A_all = [] B_all = [] u_all = [] g_all = [] sigma1_all = [] sigma2_all = [] sigma3_all = [] sigma4_all = [] l_a_posteriori1 = [] l_a_posteriori2 = [] l_a_posteriori3 = [] l_a_posteriori4 = [] P_a_posteriori1 = [] P_a_posteriori2 = [] P_a_posteriori3 = [] P_a_posteriori4 = [] env_states1 = [] env_states2 = [] env_states3 = [] env_states4 = [] KF_preds1 = [] KF_preds2 = [] KF_preds3 = [] KF_preds4 = [] squared_error1 = [] squared_error2 = [] squared_error3 = [] squared_error4 = [] KF1_params = [l_a_posteriori1,P_a_posteriori1,env_states1, KF_preds1, squared_error1] KF2_params = [l_a_posteriori2,P_a_posteriori2,env_states2, KF_preds2, squared_error2] KF3_params = [l_a_posteriori3,P_a_posteriori3,env_states3, KF_preds3, squared_error3] KF4_params = [l_a_posteriori4,P_a_posteriori4,env_states4, KF_preds4, squared_error4] '''Prediction function variables''' look_ahead_preds = [] LA_A_all = [] LA_B_all = [] LA_u_all = [] LA_g_all = [] LA_sigma1_all = [] LA_sigma2_all = [] LA_sigma3_all = [] LA_sigma4_all = [] LA_l_a_posteriori1 = [] LA_l_a_posteriori2 = [] LA_l_a_posteriori3 = [] LA_l_a_posteriori4 = [] LA_P_a_posteriori1 = [] LA_P_a_posteriori2 = [] LA_P_a_posteriori3 = [] LA_P_a_posteriori4 = [] look_ahead_vars = [look_ahead_preds,LA_A_all,LA_B_all,LA_u_all,LA_g_all,LA_sigma1_all,LA_sigma2_all,LA_sigma3_all,LA_sigma4_all, LA_l_a_posteriori1,LA_l_a_posteriori2,LA_l_a_posteriori3,LA_l_a_posteriori4, LA_P_a_posteriori1,LA_P_a_posteriori2,LA_P_a_posteriori3,LA_P_a_posteriori4] '''p-quantile loss''' Q50_numerator = np.zeros(4) Q90_numerator = np.zeros(4) '''Build LSTM''' # self.build_LSTM() '''Start gym environment''' observation=self.env.reset() '''Get initial lstm state and input, get first output/state''' initial_state = self.cell.get_initial_state(batch_size=1,dtype = tf.float64) initial_input = tf.concat((self.env_params,tf.expand_dims(tf.convert_to_tensor(observation,dtype=tf.float64),0)), axis=1) output_single, state_single = self.cell(inputs=initial_input, state=initial_state) # if not self.control or self.control_type =='uniform random': # self.variables.extend(self.cell.trainable_variables) # print('LSTM cell trainable',len(self.cell.trainable_variables)) # print('Rewards', self.rewards) # print('VARIABLES',[x.name for x in self.cell.trainable_variables]) # print('\n\n\nWEIGHTS',[x.name for x in self.cell.trainable_weights]) '''Calculate mu_0,Sigma_0, distribution using initial LSTM output''' container = tf.contrib.eager.EagerVariableStore() # control_container = tf.contrib.eager.EagerVariableStore() # with container.as_default(): # mu_0 = tf.layers.Dense(output_single, self.m, kernel_regularizer = reg.l2(self.weight_beta), # bias_regularizer = reg.l2(self.bias_beta), # name = 'mu_0dense', reuse = True) # Sigma_0 = tf.layers.Dense(output_single, self.m, kernel_regularizer = reg.l2(self.weight_beta), # bias_regularizer = reg.l2(self.bias_beta), # name = 'Sigma_0dense', reuse = True) mu_0 = self.mu_0_NN(output_single) Sigma_0 = self.Sigma_0_NN(output_single) mu_0 = tf.reshape(mu_0, shape = (self.m,1)) mu_0 = ((self.mu_0_upper_bound-self.mu_0_lower_bound)/(1+tf.exp(-mu_0)))+self.mu_0_lower_bound Sigma_0 = tf.reshape(Sigma_0, shape = (self.m,1)) Sigma_0 = tf.matmul(Sigma_0,Sigma_0,transpose_b=True)+tf.eye(4, dtype=tf.float64)*1e-8 '''Calculate predicted initial distribution''' l_0_dist = tfd.MultivariateNormalFullCovariance(loc = tf.squeeze(mu_0), covariance_matrix= Sigma_0, validate_args=True) l_0 = tf.expand_dims(l_0_dist.sample(),1) l_a_posteriori1.append(l_0) l_a_posteriori2.append(l_0) l_a_posteriori3.append(l_0) l_a_posteriori4.append(l_0) P_a_posteriori1.append(tf.eye(4, dtype = tf.float64)*self.initial_variance_estimate) P_a_posteriori2.append(tf.eye(4, dtype = tf.float64)*self.initial_variance_estimate) P_a_posteriori3.append(tf.eye(4, dtype = tf.float64)*self.initial_variance_estimate) P_a_posteriori4.append(tf.eye(4, dtype = tf.float64)*self.initial_variance_estimate) LA_l_a_posteriori1.append(l_0) LA_l_a_posteriori2.append(l_0) LA_l_a_posteriori3.append(l_0) LA_l_a_posteriori4.append(l_0) LA_P_a_posteriori1.append(tf.eye(4, dtype = tf.float64)*self.initial_variance_estimate) LA_P_a_posteriori2.append(tf.eye(4, dtype = tf.float64)*self.initial_variance_estimate) LA_P_a_posteriori3.append(tf.eye(4, dtype = tf.float64)*self.initial_variance_estimate) LA_P_a_posteriori4.append(tf.eye(4, dtype = tf.float64)*self.initial_variance_estimate) '''set initial u for random uniform control''' u = tf.Variable([0.0], dtype = tf.float64, trainable = False) first_pass = True done = False while not done: if view and self.control: self.env.render() # with container.as_default(): A, g, sigma1, sigma2, sigma3, sigma4 = self.step(output_single) if self.control: B, u = self.control_step(output_single, u, A, observation) else: B = tf.zeros(shape = (self.m,self.r), dtype = tf.float64) u = tf.zeros(shape = [1,self.r], dtype=tf.float64) '''If this is first pass in loop, add variables to graph''' # if first_pass: # self.variables.extend(container.trainable_variables()) # first_pass = False observation, reward, done, info = self.env.step(tf.squeeze(u)) '''Calculate: A,B,u,g,C,sigma,l_a_posteriori,P_a_posteriori,env_states''' KF1_update = split_forward_filter_fn(A,B,u,g,self.C_1,sigma1,l_a_posteriori1[-1],P_a_posteriori1[-1], tf.convert_to_tensor(observation[0],dtype=tf.float64)) KF2_update = split_forward_filter_fn(A,B,u,g,self.C_2,sigma2,l_a_posteriori2[-1],P_a_posteriori2[-1], tf.convert_to_tensor(observation[1],dtype=tf.float64)) KF3_update = split_forward_filter_fn(A,B,u,g,self.C_3,sigma3,l_a_posteriori3[-1],P_a_posteriori3[-1], tf.convert_to_tensor(observation[2],dtype=tf.float64)) KF4_update = split_forward_filter_fn(A,B,u,g,self.C_4,sigma4,l_a_posteriori4[-1],P_a_posteriori4[-1], tf.convert_to_tensor(observation[3],dtype=tf.float64)) '''Update lists: A_all,B_all,u_all,g_all,C_all,sigma_all,l_a_posteriori,P_a_posteriori,env_states''' A_all.append(A) B_all.append(B) u_all.append(u) g_all.append(g) sigma1_all.append(sigma1) sigma2_all.append(sigma2) sigma3_all.append(sigma3) sigma4_all.append(sigma4) for KF_single,KF_param in zip(KF1_update,KF1_params): KF_param.append(KF_single) for KF_single,KF_param in zip(KF2_update,KF2_params): KF_param.append(KF_single) for KF_single,KF_param in zip(KF3_update,KF3_params): KF_param.append(KF_single) for KF_single,KF_param in zip(KF4_update,KF4_params): KF_param.append(KF_single) next_input = tf.concat((self.env_params,env_states1[-1],env_states2[-1], env_states3[-1],env_states4[-1]),axis=1) output_single,state_single=self.cell(inputs=next_input,state=state_single) if rewards%prediction_horizon==0: # LA_preds,LA_A,LA_B,LA_u,LA_g,LA_sigma1,LA_sigma2,LA_sigma3,LA_sigma4,LA_l_a_posteriori, LA_P_a_posteriori LA_update = self.look_ahead_prediction(prediction_horizon, observation, output_single, state_single, l_a_posteriori1[-1],l_a_posteriori2[-1],l_a_posteriori3[-1],l_a_posteriori4[-1], P_a_posteriori1[-1],P_a_posteriori2[-1],P_a_posteriori3[-1],P_a_posteriori4[-1]) for var,param in zip(look_ahead_vars,LA_update): var.extend(param) # look_ahead_preds.extend(LA_preds) # LA_A_all.extend(LA_A) # LA_B_all.extend(LA_B) # LA_u_all.extend(LA_u) # LA_g_all.extend(LA_g) # LA_sigma1_all.extend(LA_sigma1) # LA_sigma2_all.extend(LA_sigma2) # LA_sigma3_all.extend(LA_sigma3) # LA_sigma4_all.extend(LA_sigma4) # LA_l_a_posteriori1.extend(LA_l_a_posteriori[0]) # LA_l_a_posteriori2.extend(LA_l_a_posteriori[1]) # LA_l_a_posteriori3.extend(LA_l_a_posteriori[2]) # LA_l_a_posteriori4.extend(LA_l_a_posteriori[3]) # LA_P_a_posteriori1.extend(LA_P_a_posteriori[0]) # LA_P_a_posteriori2.extend(LA_P_a_posteriori[1]) # LA_P_a_posteriori3.extend(LA_P_a_posteriori[2]) # LA_P_a_posteriori4.extend(LA_P_a_posteriori[3]) rewards+=1 LA_A_all = LA_A_all[:rewards] LA_B_all = LA_B_all[:rewards] LA_u_all = LA_u_all[:rewards] LA_g_all = LA_g_all[:rewards] LA_sigma1_all = LA_sigma1_all[:rewards] LA_sigma2_all = LA_sigma2_all[:rewards] LA_sigma3_all = LA_sigma3_all[:rewards] LA_sigma4_all = LA_sigma4_all[:rewards] LA_l_a_posteriori1 = LA_l_a_posteriori1[:rewards+1] LA_l_a_posteriori2 = LA_l_a_posteriori2[:rewards+1] LA_l_a_posteriori3 = LA_l_a_posteriori3[:rewards+1] LA_l_a_posteriori4 = LA_l_a_posteriori4[:rewards+1] LA_P_a_posteriori1 = LA_P_a_posteriori1[:rewards+1] LA_P_a_posteriori2 = LA_P_a_posteriori2[:rewards+1] LA_P_a_posteriori3 = LA_P_a_posteriori3[:rewards+1] LA_P_a_posteriori4 = LA_P_a_posteriori4[:rewards+1] if view and self.control: self.env.close() # param_names = ['A_all','B_all','u_all','g_all','C_all','sigma_all', # 'l_a_posteriori','P_a_posteriori','env_states','preds'] # for name,KF_param in zip(param_names,all_KF_params): # print(name,len(KF_param), KF_param[0].shape) # print('mu_0',mu_0) # print('Sigma_0',Sigma_0) # print('A_all',A_all[0]) # print('B_all',B_all[0]) # print('u_all',u_all[0]) # print('C_1',C_1) # print('sigma1_all',sigma1_all[0]) '''Start Maximum a posteriori section''' mu_11 = tf.add(tf.matmul(tf.slice(A_all[0],(0,0),(1,4)), mu_0),tf.matmul(tf.slice(B_all[0],(0,0),(1,1)),u_all[0])) Sigma_11 = tf.add(tf.matmul(tf.matmul(self.C_1,Sigma_0),self.C_1, transpose_b=True), tf.matmul(sigma1_all[0],sigma1_all[0],transpose_b=True)) mu_12 = tf.add(tf.matmul(tf.slice(A_all[0],(1,0),(1,4)), mu_0),tf.matmul(tf.slice(B_all[0],(1,0),(1,1)),u_all[0])) Sigma_12 = tf.add(tf.matmul(tf.matmul(self.C_2,Sigma_0),self.C_2, transpose_b=True), tf.matmul(sigma2_all[0],sigma2_all[0],transpose_b=True)) mu_13 = tf.add(tf.matmul(tf.slice(A_all[0],(2,0),(1,4)), mu_0),tf.matmul(tf.slice(B_all[0],(2,0),(1,1)),u_all[0])) Sigma_13 = tf.add(tf.matmul(tf.matmul(self.C_3,Sigma_0),self.C_3, transpose_b=True), tf.matmul(sigma3_all[0],sigma3_all[0],transpose_b=True)) mu_14 = tf.add(tf.matmul(tf.slice(A_all[0],(3,0),(1,4)), mu_0),tf.matmul(tf.slice(B_all[0],(3,0),(1,1)),u_all[0])) Sigma_14 = tf.add(tf.matmul(tf.matmul(self.C_4,Sigma_0),self.C_4, transpose_b=True), tf.matmul(sigma4_all[0],sigma4_all[0],transpose_b=True)) LA_mu_11 = tf.add(tf.matmul(tf.slice(LA_A_all[0],(0,0),(1,4)), mu_0),tf.matmul(tf.slice(LA_B_all[0],(0,0),(1,1)),LA_u_all[0])) LA_Sigma_11 = tf.add(tf.matmul(tf.matmul(self.C_1,Sigma_0),self.C_1, transpose_b=True), tf.matmul(LA_sigma1_all[0],LA_sigma1_all[0],transpose_b=True)) LA_mu_12 = tf.add(tf.matmul(tf.slice(LA_A_all[0],(1,0),(1,4)), mu_0),tf.matmul(tf.slice(LA_B_all[0],(1,0),(1,1)),LA_u_all[0])) LA_Sigma_12 = tf.add(tf.matmul(tf.matmul(self.C_2,Sigma_0),self.C_2, transpose_b=True), tf.matmul(LA_sigma2_all[0],LA_sigma2_all[0],transpose_b=True)) LA_mu_13 = tf.add(tf.matmul(tf.slice(LA_A_all[0],(2,0),(1,4)), mu_0),tf.matmul(tf.slice(LA_B_all[0],(2,0),(1,1)),LA_u_all[0])) LA_Sigma_13 = tf.add(tf.matmul(tf.matmul(self.C_3,Sigma_0),self.C_3, transpose_b=True), tf.matmul(LA_sigma3_all[0],LA_sigma3_all[0],transpose_b=True)) LA_mu_14 = tf.add(tf.matmul(tf.slice(LA_A_all[0],(3,0),(1,4)), mu_0),tf.matmul(tf.slice(LA_B_all[0],(3,0),(1,1)),LA_u_all[0])) LA_Sigma_14 = tf.add(tf.matmul(tf.matmul(self.C_4,Sigma_0),self.C_4, transpose_b=True), tf.matmul(LA_sigma4_all[0],LA_sigma4_all[0],transpose_b=True)) if rewards > 1: mu1,Sigma1 = self.likelihood_fn((mu_11,Sigma_11),(A_all[1:],B_all[1:],u_all[1:],g_all[1:], self.C_1,sigma1_all[1:], l_a_posteriori1[1:-1], P_a_posteriori1[1:-1])) mu2,Sigma2 = self.likelihood_fn((mu_12,Sigma_12),(A_all[1:],B_all[1:],u_all[1:],g_all[1:], self.C_2,sigma2_all[1:], l_a_posteriori2[1:-1], P_a_posteriori2[1:-1])) mu3,Sigma3 = self.likelihood_fn((mu_13,Sigma_13),(A_all[1:],B_all[1:],u_all[1:],g_all[1:], self.C_3,sigma3_all[1:], l_a_posteriori3[1:-1], P_a_posteriori3[1:-1])) mu4,Sigma4 = self.likelihood_fn((mu_14,Sigma_14),(A_all[1:],B_all[1:],u_all[1:],g_all[1:], self.C_4,sigma4_all[1:], l_a_posteriori4[1:-1], P_a_posteriori4[1:-1])) LA_mu1,LA_Sigma1 = self.likelihood_fn((LA_mu_11,LA_Sigma_11),(LA_A_all[1:],LA_B_all[1:],LA_u_all[1:],LA_g_all[1:], self.C_1,LA_sigma1_all[1:], LA_l_a_posteriori1[1:-1], LA_P_a_posteriori1[1:-1])) LA_mu2,LA_Sigma2 = self.likelihood_fn((LA_mu_12,LA_Sigma_12),(LA_A_all[1:],LA_B_all[1:],LA_u_all[1:],LA_g_all[1:], self.C_2,LA_sigma2_all[1:], LA_l_a_posteriori2[1:-1], LA_P_a_posteriori2[1:-1])) LA_mu3,LA_Sigma3 = self.likelihood_fn((LA_mu_13,LA_Sigma_13),(LA_A_all[1:],LA_B_all[1:],LA_u_all[1:],LA_g_all[1:], self.C_3,LA_sigma3_all[1:], LA_l_a_posteriori3[1:-1], LA_P_a_posteriori3[1:-1])) LA_mu4,LA_Sigma4 = self.likelihood_fn((LA_mu_14,LA_Sigma_14),(LA_A_all[1:],LA_B_all[1:],LA_u_all[1:],LA_g_all[1:], self.C_4,LA_sigma4_all[1:], LA_l_a_posteriori4[1:-1], LA_P_a_posteriori4[1:-1])) else: mu1,Sigma1 = LA_mu1,LA_Sigma1 = mu_11,Sigma_11 mu2,Sigma2 = LA_mu2,LA_Sigma2 = mu_12,Sigma_12 mu3,Sigma3 = LA_mu3,LA_Sigma3 = mu_13,Sigma_13 mu4,Sigma4 = LA_mu4,LA_Sigma4 = mu_14,Sigma_14 '''End Maximum A posteriori section''' '''p-quantile loss''' for j in range(rewards): # Q50_numerator[0] += QL(0.5,look_ahead_preds[j][0],env_states1[j]) # Q90_numerator[0] += QL(0.9,look_ahead_preds[j][0],env_states1[j]) Q50_numerator[0] += QL(0.5, KF_preds1[j], env_states1[j]) Q90_numerator[0] += QL(0.9, KF_preds1[j], env_states1[j]) for j in range(rewards): # Q50_numerator[1] += QL(0.5,look_ahead_preds[j][1],env_states2[j]) # Q90_numerator[1] += QL(0.9,look_ahead_preds[j][1],env_states2[j]) Q50_numerator[1] += QL(0.5, KF_preds2[j], env_states2[j]) Q90_numerator[1] += QL(0.9, KF_preds2[j], env_states2[j]) for j in range(rewards): # Q50_numerator[2] += QL(0.5,look_ahead_preds[j][2],env_states3[j]) # Q90_numerator[2] += QL(0.9,look_ahead_preds[j][2],env_states3[j]) Q50_numerator[2] += QL(0.5, KF_preds3[j], env_states3[j]) Q90_numerator[2] += QL(0.9, KF_preds3[j], env_states3[j]) for j in range(rewards): # Q50_numerator[3] += QL(0.5,look_ahead_preds[j][3],env_states4[j]) # Q90_numerator[3] += QL(0.9,look_ahead_preds[j][3],env_states4[j]) Q50_numerator[3] += QL(0.5, KF_preds4[j], env_states4[j]) Q90_numerator[3] += QL(0.9, KF_preds4[j], env_states4[j]) Q_denomenator1 = np.sum(np.abs(np.squeeze(np.array(env_states1))), axis = 0) Q_denomenator2 = np.sum(np.abs(np.squeeze(np.array(env_states2))), axis = 0) Q_denomenator3 = np.sum(np.abs(np.squeeze(np.array(env_states3))), axis = 0) Q_denomenator4 = np.sum(np.abs(np.squeeze(np.array(env_states4))), axis = 0) pq50_loss1 = 2*np.divide(Q50_numerator[0],Q_denomenator1) pq90_loss1 = 2*np.divide(Q90_numerator[0],Q_denomenator1) pq50_loss2 = 2*np.divide(Q50_numerator[1],Q_denomenator2) pq90_loss2 = 2*np.divide(Q90_numerator[1],Q_denomenator2) pq50_loss3 = 2*np.divide(Q50_numerator[2],Q_denomenator3) pq90_loss3 = 2*np.divide(Q90_numerator[2],Q_denomenator3) pq50_loss4 = 2*np.divide(Q50_numerator[3],Q_denomenator4) pq90_loss4 = 2*np.divide(Q90_numerator[3],Q_denomenator4) '''Compute Likelihood of observations given KF evaluation''' z1_distribution = tfd.Normal(loc = mu1, scale = Sigma1) z1_likelihood = z1_distribution.log_prob(env_states1) z2_distribution = tfd.Normal(loc = mu2, scale = Sigma2) z2_likelihood = z2_distribution.log_prob(env_states2) z3_distribution = tfd.Normal(loc = mu3, scale = Sigma3) z3_likelihood = z3_distribution.log_prob(env_states3) z4_distribution = tfd.Normal(loc = mu4, scale = Sigma4) z4_likelihood = z4_distribution.log_prob(env_states4) LA_z1_distribution = tfd.Normal(loc = LA_mu1, scale = LA_Sigma1) LA_z1_likelihood = LA_z1_distribution.log_prob(env_states1) LA_z2_distribution = tfd.Normal(loc = LA_mu2, scale = LA_Sigma2) LA_z2_likelihood = LA_z2_distribution.log_prob(env_states2) LA_z3_distribution = tfd.Normal(loc = LA_mu3, scale = LA_Sigma3) LA_z3_likelihood = LA_z3_distribution.log_prob(env_states3) LA_z4_distribution = tfd.Normal(loc = LA_mu4, scale = LA_Sigma4) LA_z4_likelihood = LA_z4_distribution.log_prob(env_states4) self.global_epoch += 1 # print('container len', len(container.variables())) # for var in container.variables(): # print(var.name) return((z1_likelihood,z2_likelihood,z3_likelihood,z4_likelihood),(LA_z1_likelihood,LA_z2_likelihood,LA_z3_likelihood,LA_z4_likelihood), rewards,tf.squeeze(tf.convert_to_tensor((KF_preds1,KF_preds2,KF_preds3,KF_preds4))), (env_states1,env_states2,env_states3,env_states4),(squared_error1,squared_error2,squared_error3,squared_error4), (pq50_loss1,pq90_loss1,pq50_loss2,pq90_loss2,pq50_loss3,pq90_loss3,pq50_loss4,pq90_loss4), look_ahead_preds)