class_docstr: A TestCase mixin for common tests on inference gym targets.
class_docstr_tok: ["A","TestCase","mixin","for","common","tests","on","inference","gym","targets","."]
docstr: Does a two-sample, two-sided Z-test between samples and ground truth.
docstr_tok: ["Does","a","two-sample",",","two-sided","Z-test","between","samples","and","ground","truth","."]
code: def _z_test(
      self,
      name,
      sample_transformation,
      transformed_samples,
      num_samples,
      mean_fudge_atol=0,
      standard_deviation_fudge_atol=0,
      sample_dims=0,
  ):
      """Does a two-sample, two-sided Z-test between samples and ground truth."""
      sample_mean = tf.nest.map_structure(
          lambda transformed_samples: np.mean(  # pylint: disable=g-long-lambda
              transformed_samples, axis=sample_dims
          ),
          transformed_samples,
      )
      sample_variance = tf.nest.map_structure(
          lambda transformed_samples: np.var(  # pylint: disable=g-long-lambda
              transformed_samples, axis=sample_dims
          ),
          transformed_samples,
      )
  
      # TODO(b/144524123): As written, this does a two sided Z-test at an
      # alpha=O(1e-7). It definitely has very little power as a result.
      # Currently it also uses the sample variance to compute the Z-score. In
      # principle, we can use the ground truth variance, but it's unclear
      # whether that's appropriate. Heuristically, a typical error that HMC has
      # is getting stuck, meaning that the sample variance is too low,
      # causing the test to fail more often. HMC can also in principle
      # over-estimate the variance, but that seems less typical.
      #
      # We should re-examine the literature for Z-testing and justify these
      # choices on formal grounds.
      def _mean_assertions_part(
          ground_truth_mean,
          ground_truth_mean_standard_error,
          sample_mean,
          sample_variance,
          num_samples,
      ):
          sample_standard_error_sq = sample_variance / num_samples
          self.assertAllClose(
              ground_truth_mean,
              sample_mean,
              # TODO(b/144290399): Use the full atol vector.
              atol=(
                  np.array(
                      5.0
                      * np.sqrt(
                          sample_standard_error_sq
                          + ground_truth_mean_standard_error**2
                      )
                  ).max()
                  + mean_fudge_atol
              ),
          )
  
      ground_truth_mean_standard_error = (
          sample_transformation.ground_truth_mean_standard_error
      )
      if ground_truth_mean_standard_error is None:
          # Note that this is NOT the intended general interpretaton of SEM being
          # None, but only the interpetation we choose for these tests.
          ground_truth_mean_standard_error = tf.nest.map_structure(
              np.zeros_like, sample_transformation.ground_truth_mean
          )
  
      self.assertAllAssertsNested(
          _mean_assertions_part,
          sample_transformation.ground_truth_mean,
          ground_truth_mean_standard_error,
          sample_mean,
          sample_variance,
          num_samples,
          msg='Comparing mean of "{}"'.format(name),
      )
  
      if (
          sample_transformation.ground_truth_standard_deviation
          is not None
      ):
          # From https://math.stackexchange.com/q/72975
          fourth_moment = tf.nest.map_structure(
              lambda transformed_samples, sample_mean: np.mean(  # pylint: disable=g-long-lambda
                  (transformed_samples - sample_mean) ** 4,
                  axis=tuple(tf.nest.flatten(sample_dims)),
              ),
              transformed_samples,
              sample_mean,
          )
  
          ground_truth_standard_deviation_standard_error = (
              sample_transformation.ground_truth_standard_deviation_standard_error
          )
          if ground_truth_standard_deviation_standard_error is None:
              # Note that this is NOT the intended general interpretaton of SESD being
              # None, but only the interpetation we choose for these tests.
              ground_truth_standard_deviation_standard_error = tf.nest.map_structure(
                  np.zeros_like,
                  sample_transformation.ground_truth_standard_deviation,
              )
  
          def _var_assertions_part(
              ground_truth_standard_deviation,
              ground_truth_standard_deviation_standard_error,
              sample_variance,
              fourth_moment,
              num_samples,
          ):
              sample_standard_error_sq = (
                  fourth_moment / num_samples
                  - sample_variance**2
                  * (num_samples - 3)
                  / num_samples
                  / (num_samples - 1)
              )
              self.assertAllClose(
                  np.square(ground_truth_standard_deviation),
                  sample_variance,
                  # TODO(b/144290399): Use the full atol vector.
                  atol=np.array(
                      5.0
                      * np.sqrt(
                          sample_standard_error_sq
                          + ground_truth_standard_deviation_standard_error
                      )
                  ).max()
                  + standard_deviation_fudge_atol,
              )
  
          self.assertAllAssertsNested(
              _var_assertions_part,
              sample_transformation.ground_truth_standard_deviation,
              ground_truth_standard_deviation_standard_error,
              sample_variance,
              fourth_moment,
              num_samples,
              msg='Comparing variance of "{}"'.format(name),
          )
  
imports: import collections
  import importlib
  from absl import flags
  from absl import logging
  import jax
  import numpy
  import tensorflow.compat.v2
  import tensorflow_probability
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 4 for calling the InferenceGymTestCase obj's underlying member method code with example usage: obj._z_test(name, sample_transformation, transformed_samples, num_samples, mean_fudge_atol, standard_deviation_fudge_atol, sample_dims) without return types