code: def testBasicSeasonality(
      self,
      constrain_mean_effect_to_zero,
      batch_shape=(),
      num_steps_per_season=None,
      num_observations_of_transition_noise=1000,
  ):
      if not ENABLE_DRIFT_SCALE_TESTS:
          print("Aborting drift scale test.")
          return
  
      # Use enough samples so we can get a reasonable estimate of the
      # distribution.
      num_samples = 1000
      distribution_batch_shape = batch_shape + (num_samples,)
  
      if num_steps_per_season is None:
          num_steps_per_season = [1, 1, 1]
      inverse_gamma_prior = inverse_gamma.InverseGamma(
          tf.constant(16.0, shape=distribution_batch_shape),
          tf.constant(4.0, shape=distribution_batch_shape),
      )
      seasonal_component = Seasonal(
          num_seasons=3,
          constrain_mean_effect_to_zero=constrain_mean_effect_to_zero,
          num_steps_per_season=num_steps_per_season,
          drift_scale_prior=transformed_distribution.TransformedDistribution(
              bijector=invert.Invert(square.Square()),
              distribution=inverse_gamma_prior,
          ),
      )
      # Scale the number of time steps to be a constant number of season
      # transitions. This means all variations should converge similarly.
      num_timesteps = num_observations_of_transition_noise * math.ceil(
          sum(num_steps_per_season) / len(num_steps_per_season)
      )
      initial_state_size = (
          seasonal_component.num_seasons - 1
          if constrain_mean_effect_to_zero
          else seasonal_component.num_seasons
      )
      drift_scale = 20.0
      ssm = seasonal_component.make_state_space_model(
          num_timesteps=num_timesteps,
          initial_state_prior=mvn_diag.MultivariateNormalDiag(
              scale_diag=[1.0] * initial_state_size
          ),
          param_vals={},
          drift_scale=drift_scale,
          # This needs to be non-zero for steps_per_season to be any value other
          # than 1 (otherwise NaN is returned after cholesky decomposition
          # failure). When there is not a season transition, LGSSM will encounter
          # variances and covariances of 0 - and even though there are no
          # differences to explain for our data (since it is generated from the
          # same LGSSM), these 0's are not supported.
          observation_noise_scale=1e-4,
      )
      seed = test_util.test_seed(sampler_type="stateless")
  
      @tf.function
      def sample_drift():
          sample = ssm.sample(sample_shape=(num_samples,), seed=seed)
          latents = ssm.posterior_sample(x=sample, seed=seed)
          return seasonal_component.experimental_resample_drift_scale(
              latents, seed=seed
          )
  
      drift_scale_samples = sample_drift()
      drift_variance_sample_mean = tf.math.reduce_mean(
          tf.math.square(drift_scale_samples)
      )
  
      # Verify that the samples of the drift scale are close to the actual drift
      # scale. This computes, given a certain number of samples, how close
      # random samples would be (a very limited duplication of what is being
      # tested).
      (
          expected_scale_variance_mean,
          expected_variance_standard_error,
      ) = _compute_summary_for_normal_samples_with_inverse_gamma_scale_prior(
          inverse_gamma_prior,
          batch_shape=batch_shape,
          num_samples=num_samples,
          num_observations_per_sample=num_observations_of_transition_noise,
          sampled_fixed_scale=drift_scale,
          seed=seed,
      )
  
      z_scores = (
          expected_scale_variance_mean - drift_variance_sample_mean
      ) / expected_variance_standard_error
      self.assertAllLess(z_scores, 4.2)
  
imports: import math
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import identity
  from tensorflow_probability.python.bijectors import invert
  from tensorflow_probability.python.bijectors import square
  from tensorflow_probability.python.distributions import categorical
  from tensorflow_probability.python.distributions import inverse_gamma
  from tensorflow_probability.python.distributions import linear_gaussian_ssm
  from tensorflow_probability.python.distributions import lognormal
  from tensorflow_probability.python.distributions import mixture_same_family
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.sts.components.seasonal import ConstrainedSeasonalStateSpaceModel
  from tensorflow_probability.python.sts.components.seasonal import Seasonal
  from tensorflow_probability.python.sts.components.seasonal import SeasonalStateSpaceModel
label_desc: Use image node_id 1 for calling the SeasonalComponentResampleDriftScaleNumerical obj's underlying member method code with example usage: obj.testBasicSeasonality(constrain_mean_effect_to_zero, batch_shape, num_steps_per_season, num_observations_of_transition_noise) and returns: seasonal_component