def test_default_construction_2d(self):
        """Tests the default parameters for 2 dimensional Brownian Motion."""
        process = BrownianMotion(dim=2)
        self.assertEqual(process.dim(), 2)
        drift_fn = process.total_drift_fn()
        # Drifts should be zero.
        t0 = np.array([0.2, 0.7, 0.9])
        delta_t = np.array([0.1, 0.8, 0.3])
        t1 = t0 + delta_t
        drifts = self.evaluate(drift_fn(t0, t1))
        self.assertEqual(drifts.shape, (3, 2))
        self.assertArrayNear(drifts.reshape([-1]), np.zeros([3 * 2]), 1e-10)
        variances = self.evaluate(process.total_covariance_fn()(t0, t1))
        self.assertEqual(variances.shape, (3, 2, 2))
        expected_variances = np.eye(2) * delta_t.reshape([-1, 1, 1])
        print(variances, expected_variances)

        self.assertArrayNear(variances.reshape([-1]),
                             expected_variances.reshape([-1]), 1e-10)
 def test_default_construction_1d(self):
     """Tests the default parameters."""
     process = BrownianMotion()
     self.assertEqual(process.dim(), 1)
     drift_fn = process.drift_fn()
     # Drift should be zero.
     t0 = np.array([0.2, 0.7, 0.9])
     t1 = t0 + [0.1, 0.8, 0.3]
     drifts = self.evaluate(drift_fn(t0, None))
     total_drift_fn = process.total_drift_fn()
     self.assertAlmostEqual(self.evaluate(total_drift_fn(0.4, 0.5)),
                            0.0,
                            places=7)
     self.assertArrayNear(drifts, np.zeros([3]), 1e-10)
     variances = self.evaluate(process.total_covariance_fn()(t0, t1))
     self.assertArrayNear(variances, t1 - t0, 1e-10)
     self.assertAlmostEqual(self.evaluate(process.total_covariance_fn()(
         0.41, 0.55)),
                            0.14,
                            places=7)
    def test_time_dependent_construction(self):
        """Tests with time dependent drift and variance."""
        def vol_fn(t):
            return tf.expand_dims(0.2 - 0.1 * tf.exp(-t), axis=-1)

        def variance_fn(t0, t1):
            # The instantaneous volatility is 0.2 - 0.1 e^(-t).
            tot_var = (t1 - t0) * 0.04 - (tf.exp(-2 * t1) -
                                          tf.exp(-2 * t0)) * 0.005
            tot_var += 0.04 * (tf.exp(-t1) - tf.exp(-t0))
            return tf.reshape(tot_var, [-1, 1, 1])

        process = BrownianMotion(dim=1,
                                 drift=0.1,
                                 volatility=vol_fn,
                                 total_covariance_fn=variance_fn)
        t0 = np.array([0.2, 0.7, 0.9])
        delta_t = np.array([0.1, 0.8, 0.3])
        t1 = t0 + delta_t
        drifts = self.evaluate(process.total_drift_fn()(t0, t1))
        self.assertArrayNear(drifts, 0.1 * delta_t, 1e-10)
        variances = self.evaluate(process.total_covariance_fn()(t0, t1))
        self.assertArrayNear(variances.reshape([-1]),
                             [0.00149104, 0.02204584, 0.00815789], 1e-8)