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)