docstr: Tests the mean and vol of the univariate GBM sampled paths.
docstr_tok: ["Tests","the","mean","and","vol","of","the","univariate","GBM","sampled","paths","."]
code: def test_univariate_sample_mean_and_variance_time_varying_drift(
      self, supply_draws, dtype
  ):
      """Tests the mean and vol of the univariate GBM sampled paths."""
      initial_state = 2.0
      min_tol = _tolerance_by_dtype(dtype)
  
      with self.subTest("Drift as a step function, sigma = 0.0"):
          mu_times = np.array([0.0, 5.0, 10.0], dtype=dtype)
          mu_values = np.array([0.0, 0.0, 0.05, 0.05], dtype=dtype)
          mu = tff.math.piecewise.PiecewiseConstantFunc(
              jump_locations=mu_times, values=mu_values, dtype=dtype
          )
          sigma = 0.0
          times = np.array([0.0, 1.0, 5.0, 7.0, 10.0], dtype=dtype)
          (
              mean,
              var,
              se_mean,
              se_var,
          ) = calculate_sample_paths_mean_and_variance(
              self,
              mu,
              sigma,
              times,
              initial_state,
              supply_draws,
              NUM_SAMPLES,
              dtype,
          )
          expected_mean = np.array(
              [
                  0.0,  # mu = 0 at t = 0
                  0.0,  # mu = 0 for t <= 1.0
                  0.0,  # mu = 0 for t < 5.0
                  2.0 * 0.05,  # mu = 0.05 for 5.0 < t <= 7.0
                  5.0 * 0.05,  # mu = 0.05 for 5.0 < t <= 10.0
              ],
              dtype=dtype,
          ) + np.log(initial_state)
          expected_var = sigma * np.sqrt(
              times
          )  # As sigma is zero this will be 0.0
          mean_tol = np.maximum(se_mean * NUM_STDERRS, min_tol)
          var_tol = np.maximum(se_var * NUM_STDERRS, min_tol)
          arrays_all_close(
              self,
              tf.squeeze(mean),
              expected_mean,
              mean_tol,
              msg="comparing means",
          )
          arrays_all_close(
              self,
              tf.squeeze(var),
              expected_var,
              var_tol,
              msg="comparing variances",
          )
  
      with self.subTest("Drift = 0.05, sigma = step function"):
          mu = 0.05
          sigma_times = np.array([0.0, 5.0, 10.0], dtype=dtype)
          sigma_values = np.array([0.0, 0.2, 0.4, 0.6], dtype=dtype)
          sigma = tff.math.piecewise.PiecewiseConstantFunc(
              jump_locations=sigma_times,
              values=sigma_values,
              dtype=dtype,
          )
          times = np.array([0.0, 1.0, 5.0, 7.0, 10.0], dtype=dtype)
          (
              mean,
              var,
              se_mean,
              se_var,
          ) = calculate_sample_paths_mean_and_variance(
              self,
              mu,
              sigma,
              times,
              initial_state,
              supply_draws,
              NUM_SAMPLES,
              dtype,
          )
          expected_mean = np.array(
              [
                  0.0,  # mu = 0 at t = 0
                  1.0 * mu - 0.5 * 1.0 * 0.2**2,  # t = 1.0
                  5.0 * mu - 0.5 * 5.0 * 0.2**2,  # t = 5.0
                  7.0 * mu
                  - 0.5 * (5.0 * 0.2**2 + 2.0 * 0.4**2),  # t = 7.0
                  10.0 * mu
                  - 0.5 * (5.0 * 0.2**2 + 5.0 * 0.4**2),  # t = 10.0
              ],
              dtype=dtype,
          ) + np.log(initial_state)
          expected_var = np.array(
              [
                  0.0,  # t = 0
                  1.0 * 0.2**2,  # t = 1.0
                  5.0 * 0.2**2,  # t = 5.0
                  5.0 * 0.2**2 + 2.0 * 0.4**2,  # t = 7.0
                  5.0 * 0.2**2 + 5.0 * 0.4**2,  # t = 10.0
              ],
              dtype=dtype,
          )
  
          # Set acceptable tolerances based on the predicted variance and a minimum
          # tolerance based on the precision.
          mean_tol = np.maximum(se_mean * NUM_STDERRS, min_tol)
          var_tol = np.maximum(se_var * NUM_STDERRS, min_tol)
  
          arrays_all_close(
              self,
              tf.squeeze(mean),
              expected_mean,
              mean_tol,
              msg="comparing means",
          )
          arrays_all_close(
              self,
              tf.squeeze(var),
              expected_var,
              var_tol,
              msg="comparing variances",
          )
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  import tf_quant_finance
  from tensorflow.python.framework import test_util
  from tf_quant_finance.models.geometric_brownian_motion import geometric_brownian_motion_test_utils
  from tf_quant_finance.models.geometric_brownian_motion import univariate_geometric_brownian_motion
label_desc: Use image node_id 12 for calling the GeometricBrownianMotionTest obj's underlying member method code with example usage: obj.test_univariate_sample_mean_and_variance_time_varying_drift(supply_draws, dtype) without return types