docstr: Tests that piecewise constant coefficients can be handled.
docstr_tok: ["Tests","that","piecewise","constant","coefficients","can","be","handled","."]
code: def test_piecewise_and_dtype(self, dtype):
      """Tests that piecewise constant coefficients can be handled."""
      mean_reversion = tff.math.piecewise.PiecewiseConstantFunc(
          jump_locations=[0.5], values=[1, 1.1], dtype=dtype
      )
      theta = tff.math.piecewise.PiecewiseConstantFunc(
          jump_locations=[0.5], values=[1, 0.9], dtype=dtype
      )
      volvol = tff.math.piecewise.PiecewiseConstantFunc(
          jump_locations=[0.3], values=[0.1, 0.2], dtype=dtype
      )
      rho = tff.math.piecewise.PiecewiseConstantFunc(
          jump_locations=[0.5], values=[0.4, 0.6], dtype=dtype
      )
      process = HestonModel(
          mean_reversion=mean_reversion,
          theta=theta,
          volvol=volvol,
          rho=rho,
          dtype=dtype,
      )
      times = [0.1, 1.0]
      num_samples = 100
      initial_state = np.array([np.log(100), 0.045], dtype=dtype)
      paths = process.sample_paths(
          times,
          time_step=0.1,
          num_samples=num_samples,
          initial_state=initial_state,
          seed=None,
      )
      paths = self.evaluate(paths)
      state_trace, volatility_trace = paths[..., 0], paths[..., 0]
      if dtype is not None:
          with self.subTest("VolatilityDtype"):
              self.assertEqual(volatility_trace.dtype, dtype)
          with self.subTest("StateDtype"):
              self.assertEqual(state_trace.dtype, dtype)
      # Check drift and volatility calculation
      dtype = paths.dtype
      initial_state = tf.convert_to_tensor(initial_state, dtype=dtype)
      with self.subTest("Drift"):
          self.assertAllClose(
              process.drift_fn()(times[0], initial_state),
              np.array([-0.0225, 0.955]),
          )
      with self.subTest("Volatility"):
          self.assertAllClose(
              process.volatility_fn()(times[0], initial_state),
              np.array([[0.21213203, 0.0], [0.00848528, 0.01944222]]),
          )
  
file_path: tf-quant-finance/tf_quant_finance/models/heston/heston_model_test.py
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v2
  import tf_quant_finance
  from tensorflow.python.framework import test_util
label_desc: Use image node_id 7 for calling the HestonModelTest obj's underlying member method code with example usage: obj.test_piecewise_and_dtype(dtype) without return types