def testSampleProbConsistentBroadcastMixNonStandardBase(self):
with self.test_session() as sess:
dims = 4
vdm = vector_diffeomixture_lib.VectorDiffeomixture(
mix_loc=[[0.], [1.]],
mix_scale=[1.],
distribution=normal_lib.Normal(1., 1.5),
loc=[
None,
np.float32([2.] * dims),
],
scale=[
linop_identity_lib.LinearOperatorScaledIdentity(
num_rows=dims,
multiplier=np.float32(1.1),
is_positive_definite=True),
linop_diag_lib.LinearOperatorDiag(
diag=np.linspace(2.5, 3.5, dims, dtype=np.float32),
is_positive_definite=True),
],
validate_args=True)
# Ball centered at component0's mean.
self.run_test_sample_consistent_log_prob(sess,
vdm,
radius=2.,
center=1.,
rtol=0.006)
# Larger ball centered at component1's mean.
self.run_test_sample_consistent_log_prob(sess,
vdm,
radius=4.,
center=3.,
rtol=0.009)
def testMeanCovarianceNoBatch(self):
with self.test_session() as sess:
dims = 3
vdm = vector_diffeomixture_lib.VectorDiffeomixture(
mix_loc=[[0.], [4.]],
mix_scale=[10.],
distribution=normal_lib.Normal(0., 1.),
loc=[
np.float32([-2.]),
None,
],
scale=[
linop_identity_lib.LinearOperatorScaledIdentity(
num_rows=dims,
multiplier=np.float32(1.5),
is_positive_definite=True),
linop_diag_lib.LinearOperatorDiag(
diag=np.linspace(2.5, 3.5, dims, dtype=np.float32),
is_positive_definite=True),
],
validate_args=True)
self.run_test_sample_consistent_mean_covariance(sess,
vdm,
rtol=0.02,
cov_rtol=0.06)
def testMeanCovarianceBatch(self):
with self.cached_session() as sess:
dims = 3
vdm = vdm_lib.VectorDiffeomixture(
mix_loc=[[0.], [4.]],
temperature=[0.1],
distribution=normal_lib.Normal(0., 1.),
loc=[
np.float32([[-2.]]),
None,
],
scale=[
linop_identity_lib.LinearOperatorScaledIdentity(
num_rows=dims,
multiplier=[np.float32(1.5)],
is_positive_definite=True),
linop_diag_lib.LinearOperatorDiag(
diag=np.stack([
np.linspace(2.5, 3.5, dims, dtype=np.float32),
np.linspace(0.5, 1.5, dims, dtype=np.float32),
]),
is_positive_definite=True),
],
quadrature_size=8,
validate_args=True)
self.run_test_sample_consistent_mean_covariance(
sess.run, vdm, rtol=0.02, cov_rtol=0.07)
def scaled_identity(w):
return linop_identity_lib.LinearOperatorScaledIdentity(
num_rows=op.range_dimension_tensor(),
multiplier=w,
is_non_singular=op.is_non_singular,
is_self_adjoint=op.is_self_adjoint,
is_positive_definite=op.is_positive_definite)
def testSampleProbConsistentBroadcastMixBatch(self):
with self.cached_session() as sess:
dims = 4
vdm = vdm_lib.VectorDiffeomixture(
mix_loc=[[0.], [1.]],
temperature=[1.],
distribution=normal_lib.Normal(0., 1.),
loc=[
None,
np.float32([2.]*dims),
],
scale=[
linop_identity_lib.LinearOperatorScaledIdentity(
num_rows=dims,
multiplier=[np.float32(1.1)],
is_positive_definite=True),
linop_diag_lib.LinearOperatorDiag(
diag=np.stack([
np.linspace(2.5, 3.5, dims, dtype=np.float32),
np.linspace(2.75, 3.25, dims, dtype=np.float32),
]),
is_positive_definite=True),
],
quadrature_size=8,
validate_args=True)
# Ball centered at component0's mean.
self.run_test_sample_consistent_log_prob(
sess.run, vdm, radius=2., center=0., rtol=0.01)
# Larger ball centered at component1's mean.
self.run_test_sample_consistent_log_prob(
sess.run, vdm, radius=4., center=2., rtol=0.01)
def testSampleProbConsistentBroadcastMixTwoBatchDims(self):
dims = 4
loc_1 = rng.randn(2, 3, dims).astype(np.float32)
with self.cached_session() as sess:
vdm = vdm_lib.VectorDiffeomixture(
mix_loc=(rng.rand(2, 3, 1) - 0.5).astype(np.float32),
temperature=[1.],
distribution=normal_lib.Normal(0., 1.),
loc=[
None,
loc_1,
],
scale=[
linop_identity_lib.LinearOperatorScaledIdentity(
num_rows=dims,
multiplier=[np.float32(1.1)],
is_positive_definite=True),
] * 2,
validate_args=True)
# Ball centered at component0's mean.
self.run_test_sample_consistent_log_prob(
sess.run, vdm, radius=2., center=0., rtol=0.01)
# Larger ball centered at component1's mean.
self.run_test_sample_consistent_log_prob(
sess.run, vdm, radius=3., center=loc_1, rtol=0.02)
def _inverse_scaled_identity(identity_operator):
return linear_operator_identity.LinearOperatorScaledIdentity(
num_rows=identity_operator._num_rows, # pylint: disable=protected-access
multiplier=1. / identity_operator.multiplier,
is_non_singular=identity_operator.is_non_singular,
is_self_adjoint=True,
is_positive_definite=identity_operator.is_positive_definite,
is_square=True)
def _cholesky_scaled_identity(identity_operator):
return linear_operator_identity.LinearOperatorScaledIdentity(
num_rows=identity_operator._num_rows, # pylint: disable=protected-access
multiplier=math_ops.sqrt(identity_operator.multiplier),
is_non_singular=True,
is_self_adjoint=True,
is_positive_definite=True,
is_square=True)
def _adjoint_scaled_identity(identity_operator):
multiplier = identity_operator.multiplier
if multiplier.dtype.is_complex:
multiplier = math_ops.conj(multiplier)
return linear_operator_identity.LinearOperatorScaledIdentity(
num_rows=identity_operator._num_rows, # pylint: disable=protected-access
multiplier=multiplier,
is_non_singular=identity_operator.is_non_singular,
is_self_adjoint=identity_operator.is_self_adjoint,
is_positive_definite=identity_operator.is_positive_definite,
is_square=True)
def _solve_linear_operator_scaled_identity(linop_a, linop_b):
"""Solve of two ScaledIdentity `LinearOperators`."""
return linear_operator_identity.LinearOperatorScaledIdentity(
num_rows=linop_a.domain_dimension_tensor(),
multiplier=linop_b.multiplier / linop_a.multiplier,
is_non_singular=registrations_util.combined_non_singular_hint(
linop_a, linop_b),
is_self_adjoint=registrations_util.
combined_commuting_self_adjoint_hint(linop_a, linop_b),
is_positive_definite=(
registrations_util.combined_commuting_positive_definite_hint(
linop_a, linop_b)),
is_square=True)
def testConcentrationLocControlsHowMuchWeightIsOnEachComponent(self):
with self.test_session() as sess:
dims = 1
vdm = vdm_lib.VectorDiffeomixture(
mix_loc=[[-1.], [0.], [1.]],
temperature=[0.5],
distribution=normal_lib.Normal(0., 1.),
loc=[
np.float32([-2.]),
np.float32([2.]),
],
scale=[
linop_identity_lib.LinearOperatorScaledIdentity(
num_rows=dims,
multiplier=np.float32(0.5),
is_positive_definite=True),
] * 2, # Use the same scale for each component.
quadrature_size=8,
validate_args=True)
samps = vdm.sample(10000)
self.assertAllEqual((10000, 3, 1), samps.shape)
samps_ = sess.run(samps).reshape(10000,
3) # Make scalar event shape.
# One characteristic of putting more weight on a component is that the
# mean is closer to that component's mean.
# Get the mean for each batch member, the names signify the value of
# concentration for that batch member.
mean_neg1, mean_0, mean_1 = samps_.mean(axis=0)
# Since concentration is the concentration for component 0,
# concentration = -1 ==> more weight on component 1, which has mean = 2
# concentration = 0 ==> equal weight
# concentration = 1 ==> more weight on component 0, which has mean = -2
self.assertLess(-2, mean_1)
self.assertLess(mean_1, mean_0)
self.assertLess(mean_0, mean_neg1)
self.assertLess(mean_neg1, 2)
# Run this test as well, just because we can.
self.run_test_sample_consistent_mean_covariance(sess.run,
vdm,
rtol=0.02,
cov_rtol=0.08)
def testSampleProbConsistentDynamicQuadrature(self):
with self.test_session() as sess:
qgrid = array_ops.placeholder(dtype=dtypes.float32)
qprobs = array_ops.placeholder(dtype=dtypes.float32)
g, p = np.polynomial.hermite.hermgauss(deg=8)
dims = 4
vdm = vector_diffeomixture_lib.VectorDiffeomixture(
mix_loc=[[0.], [1.]],
mix_scale=[1.],
distribution=normal_lib.Normal(0., 1.),
loc=[
None,
np.float32([2.] * dims),
],
scale=[
linop_identity_lib.LinearOperatorScaledIdentity(
num_rows=dims,
multiplier=np.float32(1.1),
is_positive_definite=True),
linop_diag_lib.LinearOperatorDiag(
diag=np.linspace(2.5, 3.5, dims, dtype=np.float32),
is_positive_definite=True),
],
quadrature_grid_and_probs=(g, p),
validate_args=True)
# Ball centered at component0's mean.
sess_run_fn = lambda x: sess.run(x,
feed_dict={
qgrid: g,
qprobs: p
})
self.run_test_sample_consistent_log_prob(sess_run_fn,
vdm,
radius=2.,
center=0.,
rtol=0.005)
# Larger ball centered at component1's mean.
self.run_test_sample_consistent_log_prob(sess_run_fn,
vdm,
radius=4.,
center=2.,
rtol=0.005)
def _add(self, op1, op2, operator_name, hints):
# Will build a LinearOperatorScaledIdentity.
if _type(op1) == _SCALED_IDENTITY:
multiplier_1 = op1.multiplier
else:
multiplier_1 = array_ops.ones(op1.batch_shape_tensor(), dtype=op1.dtype)
if _type(op2) == _SCALED_IDENTITY:
multiplier_2 = op2.multiplier
else:
multiplier_2 = array_ops.ones(op2.batch_shape_tensor(), dtype=op2.dtype)
return linear_operator_identity.LinearOperatorScaledIdentity(
num_rows=op1.range_dimension_tensor(),
multiplier=multiplier_1 + multiplier_2,
is_non_singular=hints.is_non_singular,
is_self_adjoint=hints.is_self_adjoint,
is_positive_definite=hints.is_positive_definite,
name=operator_name)
def testTemperatureControlsHowMuchThisLooksLikeDiscreteMixture(self):
# As temperature decreases, this should approach a mixture of normals, with
# components at -2, 2.
with self.test_session() as sess:
dims = 1
vdm = vdm_lib.VectorDiffeomixture(
mix_loc=[0.],
temperature=[[2.], [1.], [0.2]],
distribution=normal_lib.Normal(0., 1.),
loc=[
np.float32([-2.]),
np.float32([2.]),
],
scale=[
linop_identity_lib.LinearOperatorScaledIdentity(
num_rows=dims,
multiplier=np.float32(0.5),
is_positive_definite=True),
] * 2, # Use the same scale for each component.
quadrature_size=8,
validate_args=True)
samps = vdm.sample(10000)
self.assertAllEqual((10000, 3, 1), samps.shape)
samps_ = sess.run(samps).reshape(10000,
3) # Make scalar event shape.
# One characteristic of a discrete mixture (as opposed to a "smear") is
# that more weight is put near the component centers at -2, 2, and thus
# less weight is put near the origin.
prob_of_being_near_origin = (np.abs(samps_) < 1).mean(axis=0)
self.assertGreater(prob_of_being_near_origin[0],
prob_of_being_near_origin[1])
self.assertGreater(prob_of_being_near_origin[1],
prob_of_being_near_origin[2])
# Run this test as well, just because we can.
self.run_test_sample_consistent_mean_covariance(sess.run,
vdm,
rtol=0.02,
cov_rtol=0.08)
def testMeanCovarianceNoBatchUncenteredNonStandardBase(self):
with self.cached_session() as sess:
dims = 3
vdm = vdm_lib.VectorDiffeomixture(
mix_loc=[[0.], [4.]],
temperature=[0.1],
distribution=normal_lib.Normal(-1., 1.5),
loc=[
np.float32([-2.]),
np.float32([0.]),
],
scale=[
linop_identity_lib.LinearOperatorScaledIdentity(
num_rows=dims,
multiplier=np.float32(1.5),
is_positive_definite=True),
linop_diag_lib.LinearOperatorDiag(
diag=np.linspace(2.5, 3.5, dims, dtype=np.float32),
is_positive_definite=True),
],
quadrature_size=8,
validate_args=True)
self.run_test_sample_consistent_mean_covariance(
sess.run, vdm, num_samples=int(1e6), rtol=0.01, cov_atol=0.025)
def _mean_of_covariance_given_quadrature_component(self, diag_only):
p = self.mixture_distribution.probs
# To compute E[Cov(Z|V)], we'll add matrices within three categories:
# scaled-identity, diagonal, and full. Then we'll combine these at the end.
scale_identity_multiplier = None
diag = None
full = None
for k, aff in enumerate(self.interpolated_affine):
s = aff.scale # Just in case aff.scale has side-effects, we'll call once.
if (s is None or isinstance(
s, linop_identity_lib.LinearOperatorIdentity)):
scale_identity_multiplier = add(scale_identity_multiplier,
p[..., k, array_ops.newaxis])
elif isinstance(s,
linop_identity_lib.LinearOperatorScaledIdentity):
scale_identity_multiplier = add(
scale_identity_multiplier, (p[..., k, array_ops.newaxis] *
math_ops.square(s.multiplier)))
elif isinstance(s, linop_diag_lib.LinearOperatorDiag):
diag = add(diag, (p[..., k, array_ops.newaxis] *
math_ops.square(s.diag_part())))
else:
x = (p[..., k, array_ops.newaxis, array_ops.newaxis] *
s.matmul(s.to_dense(), adjoint_arg=True))
if diag_only:
x = array_ops.matrix_diag_part(x)
full = add(full, x)
# We must now account for the fact that the base distribution might have a
# non-unity variance. Recall that, since X ~ iid Law(X_0),
# `Cov(SX+m) = S Cov(X) S.T = S S.T Diag(Var(X_0))`.
# We can scale by `Var(X)` (vs `Cov(X)`) since X corresponds to `d` iid
# samples from a scalar-event distribution.
v = self.distribution.variance()
if scale_identity_multiplier is not None:
scale_identity_multiplier *= v
if diag is not None:
diag *= v[..., array_ops.newaxis]
if full is not None:
full *= v[..., array_ops.newaxis]
if diag_only:
# Apparently we don't need the full matrix, just the diagonal.
r = add(diag, full)
if r is None and scale_identity_multiplier is not None:
ones = array_ops.ones(self.event_shape_tensor(),
dtype=self.dtype)
return scale_identity_multiplier[..., array_ops.newaxis] * ones
return add(r, scale_identity_multiplier)
# `None` indicates we don't know if the result is positive-definite.
is_positive_definite = (True if all(
aff.scale.is_positive_definite
for aff in self.endpoint_affine) else None)
to_add = []
if diag is not None:
to_add.append(
linop_diag_lib.LinearOperatorDiag(
diag=diag, is_positive_definite=is_positive_definite))
if full is not None:
to_add.append(
linop_full_lib.LinearOperatorFullMatrix(
matrix=full, is_positive_definite=is_positive_definite))
if scale_identity_multiplier is not None:
to_add.append(
linop_identity_lib.LinearOperatorScaledIdentity(
num_rows=self.event_shape_tensor()[0],
multiplier=scale_identity_multiplier,
is_positive_definite=is_positive_definite))
return (linop_add_lib.add_operators(to_add)[0].to_dense()
if to_add else None)
def _inverse_block_lower_triangular(block_lower_triangular_operator):
"""Inverse of LinearOperatorBlockLowerTriangular.
We recursively apply the identity:
```none
|A 0|' = | A' 0|
|B C| |-C'BA' C'|
```
where `A` is n-by-n, `B` is m-by-n, `C` is m-by-m, and `'` denotes inverse.
This identity can be verified through multiplication:
```none
|A 0|| A' 0|
|B C||-C'BA' C'|
= | AA' 0|
|BA'-CC'BA' CC'|
= |I 0|
|0 I|
```
Args:
block_lower_triangular_operator: Instance of
`LinearOperatorBlockLowerTriangular`.
Returns:
block_lower_triangular_operator_inverse: Instance of
`LinearOperatorBlockLowerTriangular`, the inverse of
`block_lower_triangular_operator`.
"""
if len(block_lower_triangular_operator.operators) == 1:
return (
linear_operator_block_lower_triangular.
LinearOperatorBlockLowerTriangular(
[[block_lower_triangular_operator.operators[0][0].inverse()]],
is_non_singular=block_lower_triangular_operator.
is_non_singular,
is_self_adjoint=block_lower_triangular_operator.
is_self_adjoint,
is_positive_definite=(
block_lower_triangular_operator.is_positive_definite),
is_square=True))
blockwise_dim = len(block_lower_triangular_operator.operators)
# Calculate the inverse of the `LinearOperatorBlockLowerTriangular`
# representing all but the last row of `block_lower_triangular_operator` with
# a recursive call (the matrix `A'` in the docstring definition).
upper_left_inverse = (
linear_operator_block_lower_triangular.
LinearOperatorBlockLowerTriangular(
block_lower_triangular_operator.operators[:-1]).inverse())
bottom_row = block_lower_triangular_operator.operators[-1]
bottom_right_inverse = bottom_row[-1].inverse()
# Find the bottom row of the inverse (equal to `[-C'BA', C']` in the docstring
# definition, where `C` is the bottom-right operator of
# `block_lower_triangular_operator` and `B` is the set of operators in the
# bottom row excluding `C`). To find `-C'BA'`, we first iterate over the
# column partitions of `A'`.
inverse_bottom_row = []
for i in range(blockwise_dim - 1):
# Find the `i`-th block of `BA'`.
blocks = []
for j in range(i, blockwise_dim - 1):
result = bottom_row[j].matmul(upper_left_inverse.operators[j][i])
if not any(
isinstance(result, op_type) for op_type in
linear_operator_addition.SUPPORTED_OPERATORS):
result = linear_operator_full_matrix.LinearOperatorFullMatrix(
result.to_dense())
blocks.append(result)
summed_blocks = linear_operator_addition.add_operators(blocks)
assert len(summed_blocks) == 1
block = summed_blocks[0]
# Find the `i`-th block of `-C'BA'`.
block = bottom_right_inverse.matmul(block)
block = linear_operator_identity.LinearOperatorScaledIdentity(
num_rows=bottom_right_inverse.domain_dimension_tensor(),
multiplier=math_ops.cast(-1, dtype=block.dtype)).matmul(block)
inverse_bottom_row.append(block)
# `C'` is the last block of the inverted linear operator.
inverse_bottom_row.append(bottom_right_inverse)
return (
linear_operator_block_lower_triangular.
LinearOperatorBlockLowerTriangular(
upper_left_inverse.operators + [inverse_bottom_row],
is_non_singular=block_lower_triangular_operator.is_non_singular,
is_self_adjoint=block_lower_triangular_operator.is_self_adjoint,
is_positive_definite=(
block_lower_triangular_operator.is_positive_definite),
is_square=True))
