updates tests

Author: Michael Fuest

cents/data_generator.py

@@ -63,6 +63,7 @@ def __init__(
                 self.set_dataset_spec(
                     cfg.dataset, self._read_ctx_codes(cfg.dataset.name)
                 )
+                self.model_type = cfg.model.name
         elif model_name is not None:
             self.model_type = get_model_type_from_hf_name(model_name)
             self.cfg = cfg or self._default_cfg()

tests/evaluator/test_evaluator.py

@@ -59,6 +59,7 @@ def cfg(tmp_path):
                 "eval_vis": True,
                 "eval_pv_shift": False,
                 "eval_context_sparse": False,
+                "eval_disentanglement": False,
             },
         }
     )
@@ -81,7 +82,7 @@ def test_compute_metrics_sets_all_keys(evaluator):
             "context": [0, 1, 0],
         }
     )
-    evaluator.compute_metrics(real, syn, df)
+    evaluator.compute_quality_metrics(real, syn, df)
     for key in ("DTW", "MMD", "Context_FID", "Disc_Score", "Pred_Score"):
         assert key in evaluator.current_results["metrics"], f"{key} missing"
 
@@ -103,6 +104,9 @@ class DummyModel:
         def to(self, device):
             return self
 
+        def eval(self):
+            return self
+
         def generate(self, ctx):
             batch_size = next(iter(ctx.values())).shape[0]
             return torch.zeros((batch_size, 1, 1))

tests/model/test_models.py

@@ -12,7 +12,7 @@ def test_trainer_fit_diffusion_1d(dummy_trainer_diffusion_1d):
 def test_trainer_get_data_generator_diffusion_1d(dummy_trainer_diffusion_1d):
     trainer = dummy_trainer_diffusion_1d
     dg = trainer.get_data_generator()
-    assert dg.model_name == "diffusion_ts"
+    assert dg.model_type == "diffusion_ts"
     assert dg.model is not None
 
 
@@ -25,7 +25,7 @@ def test_trainer_fit_acgan_1d(dummy_trainer_acgan_1d):
 def test_trainer_get_data_generator_acgan_1d(dummy_trainer_acgan_1d):
     trainer = dummy_trainer_acgan_1d
     dg = trainer.get_data_generator()
-    assert dg.model_name == "acgan"
+    assert dg.model_type == "acgan"
     assert dg.model is not None
 
 
@@ -38,7 +38,7 @@ def test_trainer_fit_diffusion_2d(dummy_trainer_diffusion_2d):
 def test_trainer_get_data_generator_diffusion_2d(dummy_trainer_diffusion_2d):
     trainer = dummy_trainer_diffusion_2d
     dg = trainer.get_data_generator()
-    assert dg.model_name == "diffusion_ts"
+    assert dg.model_type == "diffusion_ts"
     assert dg.model is not None
 
 
@@ -51,5 +51,5 @@ def test_trainer_fit_acgan_2d(dummy_trainer_acgan_2d):
 def test_trainer_get_data_generator_acgan_2d(dummy_trainer_acgan_2d):
     trainer = dummy_trainer_acgan_2d
     dg = trainer.get_data_generator()
-    assert dg.model_name == "acgan"
+    assert dg.model_type == "acgan"
     assert dg.model is not None

tests/test_configs/evaluator/default.yaml

@@ -1,5 +1,6 @@
 model_name: ${model.name}
 eval_pv_shift: False
+eval_disentanglement: True
 eval_metrics: True
 eval_vis: True
 eval_context_sparse: False