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