# Path: disentanglement_lib/data/ground_truth/dummy_data.py
# class IdentityObservationsData(ground_truth_data.GroundTruthData):
# class DummyData(ground_truth_data.GroundTruthData):
#   def num_factors(self):
#   def observation_shape(self):
#   def factors_num_values(self):
#   def sample_factors(self, num, random_state):
#   def sample_observations_from_factors(self, factors, random_state):
#   def factor_names(self):
#   def num_factors(self):
#   def factors_num_values(self):
#   def observation_shape(self):
#   def sample_factors(self, num, random_state):
#   def sample_observations_from_factors(self, factors, random_state):
# 
# Path: disentanglement_lib/evaluation/metrics/unified_scores.py
# def unified_scores(mus_train, ys_train, mus_test, ys_test, matrix_fns,
#                    artifact_dir=None, factor_names=None):
#   """Computes unified scores."""
# 
#   scores = {}
#   kws = {}
#   for matrix_fn in matrix_fns:
#     # Matrix should have shape [num_codes, num_factors].
#     matrix = matrix_fn(mus_train, ys_train, mus_test, ys_test)
#     matrix_name = matrix_fn.__name__
#     if artifact_dir is not None:
#       visualize_scores.heat_square(matrix.copy(), artifact_dir, matrix_name,
#                                    "Latent codes", "Factors of Variation",
#                                    factor_names=factor_names)
#       visualize_scores.plot_recovery_vs_independent(matrix.copy().T,
#                                                     artifact_dir,
#                                                     matrix_name+"_pr")
#       merge_points = dendrogram.dendrogram_plot(matrix.copy().T, os.path.join(
#           artifact_dir, matrix_name+"_dendrogram"), factor_names)
#       kws[matrix_name] = merge_points
#     results_dict = pr_curves_values(matrix)
#     if matrix_name in kws:
#       kws[matrix_name].update(results_dict)
#     else:
#       kws[matrix_name] = results_dict
#     for aggregation_fn in AGGREGATION_FN:
#       results_dict = aggregation_fn(matrix, ys_train)
#       kws[matrix_name].update(results_dict)
#   scores = results.namespaced_dict(scores, **kws)
#   return scores


from absl.testing import absltest
from disentanglement_lib.data.ground_truth import dummy_data
from disentanglement_lib.evaluation.metrics import unified_scores
import numpy as np
import gin.tf
    scores = unified_scores.compute_unified_scores(
        ground_truth_data, representation_function, random_state, None, 10000,
        100, matrix_fns=[unified_scores.importance_gbt_matrix])
    self.assertBetween(
        scores["importance_gbt_matrix.dci_disentanglement"], 0.9, 1.0)

  def test_duplicated_latent_space_mig(self):
    gin.bind_parameter("discretizer.discretizer_fn", _identity_discretizer)
    gin.bind_parameter("discretizer.num_bins", 10)
    ground_truth_data = dummy_data.IdentityObservationsData()
    def representation_function(x):
      x = np.array(x, dtype=np.float64)
      return np.hstack([x, x])
    random_state = np.random.RandomState(0)
    scores = unified_scores.compute_unified_scores(
        ground_truth_data, representation_function, random_state, None, 1000,
        1000, matrix_fns=[unified_scores.mutual_information_matrix])
    self.assertBetween(
        scores["mutual_information_matrix.mig"], 0.0, 0.2)
    self.assertBetween(
        scores["mutual_information_matrix.modularity"], 0.9, 1.0)

  def test_duplicated_latent_space_dci(self):
    gin.bind_parameter("discretizer.discretizer_fn", _identity_discretizer)
    gin.bind_parameter("discretizer.num_bins", 10)
    ground_truth_data = dummy_data.IdentityObservationsData()
    def representation_function(x):
      x = np.array(x, dtype=np.float64)
      return np.hstack([x, x])
    random_state = np.random.RandomState(0)

# Path: disentanglement_lib/data/ground_truth/dummy_data.py
# class IdentityObservationsData(ground_truth_data.GroundTruthData):
# class DummyData(ground_truth_data.GroundTruthData):
#   def num_factors(self):
#   def observation_shape(self):
#   def factors_num_values(self):
#   def sample_factors(self, num, random_state):
#   def sample_observations_from_factors(self, factors, random_state):
#   def factor_names(self):
#   def num_factors(self):
#   def factors_num_values(self):
#   def observation_shape(self):
#   def sample_factors(self, num, random_state):
#   def sample_observations_from_factors(self, factors, random_state):
# 
# Path: disentanglement_lib/evaluation/metrics/unified_scores.py
# def unified_scores(mus_train, ys_train, mus_test, ys_test, matrix_fns,
#                    artifact_dir=None, factor_names=None):
#   """Computes unified scores."""
# 
#   scores = {}
#   kws = {}
#   for matrix_fn in matrix_fns:
#     # Matrix should have shape [num_codes, num_factors].
#     matrix = matrix_fn(mus_train, ys_train, mus_test, ys_test)
#     matrix_name = matrix_fn.__name__
#     if artifact_dir is not None:
#       visualize_scores.heat_square(matrix.copy(), artifact_dir, matrix_name,
#                                    "Latent codes", "Factors of Variation",
#                                    factor_names=factor_names)
#       visualize_scores.plot_recovery_vs_independent(matrix.copy().T,
#                                                     artifact_dir,
#                                                     matrix_name+"_pr")
#       merge_points = dendrogram.dendrogram_plot(matrix.copy().T, os.path.join(
#           artifact_dir, matrix_name+"_dendrogram"), factor_names)
#       kws[matrix_name] = merge_points
#     results_dict = pr_curves_values(matrix)
#     if matrix_name in kws:
#       kws[matrix_name].update(results_dict)
#     else:
#       kws[matrix_name] = results_dict
#     for aggregation_fn in AGGREGATION_FN:
#       results_dict = aggregation_fn(matrix, ys_train)
#       kws[matrix_name].update(results_dict)
#   scores = results.namespaced_dict(scores, **kws)
#   return scores


# Path: disentanglement_lib/evaluation/metrics/unified_scores_test.py
from absl.testing import absltest
from disentanglement_lib.data.ground_truth import dummy_data
from disentanglement_lib.evaluation.metrics import unified_scores
import numpy as np
import gin.tf

    scores = unified_scores.compute_unified_scores(
        ground_truth_data, representation_function, random_state, None, 10000,
        100, matrix_fns=[unified_scores.importance_gbt_matrix])
    self.assertBetween(
        scores["importance_gbt_matrix.dci_disentanglement"], 0.9, 1.0)

  def test_duplicated_latent_space_mig(self):
    gin.bind_parameter("discretizer.discretizer_fn", _identity_discretizer)
    gin.bind_parameter("discretizer.num_bins", 10)
    ground_truth_data = dummy_data.IdentityObservationsData()
    def representation_function(x):
      x = np.array(x, dtype=np.float64)
      return np.hstack([x, x])
    random_state = np.random.RandomState(0)
    scores = unified_scores.compute_unified_scores(
        ground_truth_data, representation_function, random_state, None, 1000,
        1000, matrix_fns=[unified_scores.mutual_information_matrix])
    self.assertBetween(
        scores["mutual_information_matrix.mig"], 0.0, 0.2)
    self.assertBetween(
        scores["mutual_information_matrix.modularity"], 0.9, 1.0)

  def test_duplicated_latent_space_dci(self):
    gin.bind_parameter("discretizer.discretizer_fn", _identity_discretizer)
    gin.bind_parameter("discretizer.num_bins", 10)
    ground_truth_data = dummy_data.IdentityObservationsData()
    def representation_function(x):
      x = np.array(x, dtype=np.float64)
      return np.hstack([x, x])
    random_state = np.random.RandomState(0)

    scores = unified_scores.compute_unified_scores(