def from_loc_and_scale(cls,
loc,
scale,
low=0.,
high=1e10,
scale_shift=1e-7):
"""
Instantiate a learnable distribution with good default bijectors.
loc : array
The initial location of the distribution
scale : array
The initial scale parameter of the distribution
low : float or array (optional)
The lower limit of the support for the distribution.
high : float or array (optional)
The upper limit of the support for the distribution.
scale_shift : float (optional)
A small constant added to the scale to increase numerical stability.
"""
loc = tfp.util.TransformedVariable(
loc,
tfb.Softplus(),
)
scale = tfp.util.TransformedVariable(
scale,
tfb.Chain([
tfb.Softplus(),
tfb.Shift(scale_shift),
]),
)
return cls(loc, scale, low, high)
def test_stochastic_optimization(self):
seed = test_util.test_seed_stream()
tf.random.set_seed(seed())
# Example of fitting one normal to another using a
# Monte Carlo variational loss.
locs = tf.Variable(tf.random.normal([10], seed=seed()))
scales = tfp.util.TransformedVariable(
tf.nn.softplus(tf.random.normal([10], seed=seed())),
tfb.Softplus())
trained_dist = tfd.Normal(locs, scales)
target_dist = tfd.Normal(loc=-0.4, scale=1.2)
optimizer = tf.optimizers.Adam(learning_rate=0.1)
@tf.function(autograph=False)
def optimization_step():
with tf.GradientTape() as tape:
loss = tfp.vi.monte_carlo_variational_loss(
target_log_prob_fn=target_dist.log_prob,
surrogate_posterior=trained_dist,
sample_size=20,
seed=seed())
grads = tape.gradient(loss, trained_dist.trainable_variables)
optimizer.apply_gradients(
zip(grads, trained_dist.trainable_variables))
return loss, grads
criterion = (tfp.optimizer.convergence_criteria.
SuccessiveGradientsAreUncorrelated(window_size=10,
min_num_steps=20))
loss, grads = optimization_step()
self.evaluate(tf1.global_variables_initializer())
auxiliary_state = criterion.bootstrap(loss, grads,
trained_dist.trainable_variables)
for step in range(1, 100):
loss, grads = optimization_step()
has_converged, auxiliary_state = criterion.one_step(
step, loss, grads, trained_dist.trainable_variables,
auxiliary_state)
has_converged_ = self.evaluate(has_converged)
if has_converged_:
break
# Check that the criterion successfully stopped the optimization
# (at step 32 with the test seed as of this writing).
self.assertLess(step, 99) # pylint: disable=undefined-loop-variable
# Because this is a stochastic optimization with no learning rate decay,
# we will not converge to the true values, just to a stationary distribution
# that (hopefully) includes them.
self.assertLess(self.evaluate(tf.reduce_sum(loss)), 1.5)
self.assertAllClose(*self.evaluate(
(tf.reduce_mean(locs), target_dist.mean())),
atol=0.5)
self.assertAllClose(*self.evaluate(
(tf.reduce_mean(scales), target_dist.stddev())),
atol=0.5)
def __init__(self, scale=None, **kwargs):
if scale is None:
scaling = tfb.Identity()
else:
scaling = tfb.Scale(scale)
transform = tfb.Chain([tfb.Softplus(), scaling])
super().__init__(transform, **kwargs)
def test_laue(likelihood_model, prior_model, scaling_model, laue_inputs,
mc_samples):
nrefls = np.max(BaseModel.get_refl_id(laue_inputs)) + 1
n_images = np.max(BaseModel.get_image_id(laue_inputs)) + 1
#For the students
dof = 4.
if likelihood_model == StudentTLikelihood:
likelihood = likelihood_model(dof)
else:
likelihood = likelihood_model()
if prior_model == WilsonPrior:
prior = prior_model(
np.random.choice([True, False], nrefls),
np.ones(nrefls).astype('float32'),
)
elif prior_model == StudentTReferencePrior:
prior = prior_model(
np.ones(nrefls).astype('float32'),
np.ones(nrefls).astype('float32'), dof)
else:
prior = prior_model(
np.ones(nrefls).astype('float32'),
np.ones(nrefls).astype('float32'),
)
mlp_scaler = MLPScaler(2, 3)
if scaling_model == HybridImageScaler:
image_scaler = ImageScaler(n_images)
scaler = HybridImageScaler(mlp_scaler, image_scaler)
elif scaling_model == MLPScaler:
scaler = mlp_scaler
surrogate_posterior = tfd.TruncatedNormal(
tf.Variable(prior.mean()),
tfp.util.TransformedVariable(
prior.stddev() / 10.,
tfb.Softplus(),
),
low=1e-5,
high=1e10,
)
merger = VariationalMergingModel(surrogate_posterior, prior, likelihood,
scaler, mc_samples)
ipred = merger(laue_inputs)
isfinite = np.all(np.isfinite(ipred.numpy()))
assert isfinite
merger = VariationalMergingModel(surrogate_posterior, prior, likelihood,
scaler)
merger.compile('Adam')
def __init__(self, scale=None, skewness=None, tailweight=None, **kwargs):
if scale is None:
scaling = tfb.Identity()
else:
scaling = tfb.Scale(scale)
transform = tfb.Chain([
scaling,
tfb.Softplus(),
tfb.SinhArcsinh(
skewness,
tailweight,
),
])
super().__init__(transform, **kwargs)
# X represents the full set of points for which we have
# measured values
X = np.vstack((grid_x,grid_y)).T
# y is the array of values at each x
y = np.sin(grid_x)*np.sin(grid_y/10.)
if dtype == tf.float32:
X = X.astype(np.float32)
y = y.astype(np.float32)
# These are kernel parameters. We will use an automatic relevance determining rbf kernel
amplitude = tfp.util.TransformedVariable(
tf.Variable(1., dtype=dtype), tfb.Softplus(), dtype=dtype, name='amplitude')
length_scale = tfp.util.TransformedVariable(
tf.Variable(tf.ones(2, dtype=dtype)), tfb.Softplus(), dtype=dtype, name='length_scale')
#The FeatureScaled object is how you create ARD kernels in tfp
kernel = tfp.math.psd_kernels.FeatureScaled(
tfp.math.psd_kernels.ExponentiatedQuadratic(amplitude=amplitude),
length_scale
)
#This is the overall amount of noise associated with each y-value
#This is called β^{-1} in the Hensman paper. In practice, this could
#be tuned or learned.
observation_noise_variance = tfp.util.TransformedVariable(
1., tfb.Softplus(), dtype=dtype, name='observation_noise_variance')
)).reshape((2, gridpoints**2))
# y is the array of values at each x
y = np.sin(grid_x) * np.sin(grid_y / 10.)
# X represents the full set of points for which we have
# measured values
X = np.vstack((grid_x, grid_y)).T
if dtype == tf.float32:
X = X.astype(np.float32)
y = y.astype(np.float32)
# These are kernel parameters. We will use an automatic relevance determining rbf kernel
amplitude = tfp.util.TransformedVariable(tf.Variable(1., dtype=dtype),
tfb.Softplus(),
dtype=dtype,
name='amplitude')
length_scale = tfp.util.TransformedVariable(tf.Variable(tf.ones(2,
dtype=dtype)),
tfb.Softplus(),
dtype=dtype,
name='length_scale')
#The FeatureScaled object is how you create ARD kernels in tfp
kernel = tfp.math.psd_kernels.FeatureScaled(
tfp.math.psd_kernels.ExponentiatedQuadratic(amplitude=amplitude),
length_scale)
#This is the overall amount of noise associated with each y-value
#This is called β^{-1} in the Hensman paper. In practice, this could