def test_shapes(self):
"""Tests the sample shapes."""
sample_no_batch = self.evaluate(
tff_rnd.mv_normal_sample([2, 4], mean=[0.2, 0.1]))
self.assertEqual(sample_no_batch.shape, (2, 4, 2))
sample_batch = self.evaluate(
tff_rnd.mv_normal_sample([2, 4],
mean=[[0.2, 0.1], [0., -0.1], [0., 0.1]]))
self.assertEqual(sample_batch.shape, (2, 4, 3, 2))
def test_mean_and_scale(self):
"""Tests sample for scale specification."""
mean = np.array([[1.0, 0.1], [0.1, 1.0]])
scale = np.array([[0.4, -0.1], [0.22, 1.38]])
covariance = np.matmul(scale, scale.transpose())
size = 30000
sample = self.evaluate(
tff_rnd.mv_normal_sample([size],
mean=mean,
scale_matrix=scale,
seed=7534))
np.testing.assert_array_equal(sample.shape, [size, 2, 2])
np.testing.assert_array_almost_equal(np.mean(sample, axis=0),
mean,
decimal=1)
np.testing.assert_array_almost_equal(np.cov(sample[:, 0, :],
rowvar=False),
covariance,
decimal=1)
np.testing.assert_array_almost_equal(np.cov(sample[:, 1, :],
rowvar=False),
covariance,
decimal=1)
def test_general_mean_covariance(self):
"""Tests that the sample is correctly generated for general params."""
mean = np.array([[1.0, 0.1], [0.1, 1.0]])
covar = np.array([
[[0.9, -0.1], [-0.1, 1.0]],
[[1.1, -0.3], [-0.3, 0.6]],
])
size = 30000
sample = self.evaluate(
tff_rnd.mv_normal_sample([size],
mean=mean,
covariance_matrix=covar,
seed=4567))
np.testing.assert_array_equal(sample.shape, [size, 2, 2])
np.testing.assert_array_almost_equal(np.mean(sample, axis=0),
mean,
decimal=1)
np.testing.assert_array_almost_equal(np.cov(sample[:, 0, :],
rowvar=False),
covar[0],
decimal=1)
np.testing.assert_array_almost_equal(np.cov(sample[:, 1, :],
rowvar=False),
covar[1],
decimal=1)
def test_mean_default(self):
"""Tests that the default value of mean is 0."""
covar = np.array([[1.0, 0.1], [0.1, 1.0]])
sample = self.evaluate(
tff_rnd.mv_normal_sample([40000],
covariance_matrix=covar,
seed=1234))
np.testing.assert_array_equal(sample.shape, [40000, 2])
self.assertArrayNear(np.mean(sample, axis=0), [0.0, 0.0], 1e-2)
self.assertArrayNear(
np.cov(sample, rowvar=False).reshape([-1]), covar.reshape([-1]),
2e-2)
def test_antithetic_sample_requires_even_dim(self):
"""Error is trigerred if the first dim of sample_shape is odd."""
mean = np.array([[1.0, 0.1], [0.1, 1.0]])
scale = np.array([[0.4, -0.1], [0.22, 1.38]])
sample_shape = [11, 100]
# Should fail: The first dimension of `sample_shape` should be even.
with self.assertRaises(tf.errors.InvalidArgumentError):
self.evaluate(
tff_rnd.mv_normal_sample(
sample_shape,
mean=mean,
scale_matrix=scale,
random_type=tff_rnd.RandomType.PSEUDO_ANTITHETIC))
def test_covariance_default(self):
"""Tests that the default value of the covariance matrix is identity."""
mean = np.array([[1.0, 0.1], [0.1, 1.0]])
sample = self.evaluate(tff_rnd.mv_normal_sample([10000], mean=mean))
np.testing.assert_array_equal(sample.shape, [10000, 2, 2])
np.testing.assert_array_almost_equal(np.mean(sample, axis=0),
mean,
decimal=1)
np.testing.assert_array_almost_equal(np.cov(sample[:, 0, :],
rowvar=False),
np.eye(2),
decimal=1)
np.testing.assert_array_almost_equal(np.cov(sample[:, 1, :],
rowvar=False),
np.eye(2),
decimal=1)
def step_fn(i, written_count, current_state, result):
"""Performs one step of Euler scheme."""
current_time = times[i + 1]
dw = random.mv_normal_sample((num_samples,),
mean=wiener_mean,
random_type=random_type,
seed=seed)
dw = dw * sqrt_dt[i]
dt_inc = dt[i] * self.drift_fn()(current_time, current_state) # pylint: disable=not-callable
dw_inc = tf.squeeze(
tf.matmul(self.volatility_fn()(current_time, current_state), dw), -1) # pylint: disable=not-callable
next_state = current_state + dt_inc + dw_inc
# Keep only states for times, requested by user.
result = tf.cond(keep_mask[i + 1],
(lambda: result.write(written_count, next_state)),
(lambda: result))
written_count += tf.cast(keep_mask[i + 1], dtype=tf.int32)
return (i + 1, written_count, next_state, result)
def test_mean_and_scale_antithetic(self):
"""Tests antithetic sampler for scale specification."""
mean = np.array([[1.0, 0.1], [0.1, 1.0]])
scale = np.array([[0.4, -0.1], [0.22, 1.38]])
covariance = np.matmul(scale, scale.transpose())
size = 30000
sample = self.evaluate(
tff_rnd.mv_normal_sample(
[size],
mean=mean,
scale_matrix=scale,
random_type=tff_rnd.RandomType.PSEUDO_ANTITHETIC,
seed=42))
np.testing.assert_array_equal(sample.shape, [size, 2, 2])
# Antithetic combination of samples should be equal to the `mean`
antithetic_size = size // 2
antithetic_combination = (sample[:antithetic_size, ...] +
sample[antithetic_size:, ...]) / 2
np.testing.assert_allclose(antithetic_combination,
mean + np.zeros([antithetic_size, 2, 2]),
1e-10, 1e-10)
# Get the antithetic pairs and verify normality
np.testing.assert_array_almost_equal(np.mean(sample[:antithetic_size,
...],
axis=0),
mean,
decimal=1)
np.testing.assert_array_almost_equal(np.cov(sample[:antithetic_size,
0, :],
rowvar=False),
covariance,
decimal=1)
np.testing.assert_array_almost_equal(np.cov(sample[:antithetic_size,
1, :],
rowvar=False),
covariance,
decimal=1)
def _euler_step(*, i, written_count, current_state, drift_fn, volatility_fn,
wiener_mean, num_samples, times, dt, sqrt_dt, keep_mask,
random_type, seed, normal_draws, result, record_samples):
"""Performs one step of Euler scheme."""
current_time = times[i + 1]
written_count = tf.cast(written_count, tf.int32)
if normal_draws is not None:
dw = normal_draws[i]
else:
dw = random.mv_normal_sample((num_samples, ),
mean=wiener_mean,
random_type=random_type,
seed=seed)
dw = dw * sqrt_dt[i]
dt_inc = dt[i] * drift_fn(current_time, current_state) # pylint: disable=not-callable
dw_inc = tf.linalg.matvec(volatility_fn(current_time, current_state), dw) # pylint: disable=not-callable
next_state = current_state + dt_inc + dw_inc
if record_samples:
result = result.write(written_count, next_state)
else:
result = next_state
written_count += tf.cast(keep_mask[i + 1], dtype=tf.int32)
return i + 1, written_count, next_state, result
