Update callbacks and tests (#482)

---------

Co-authored-by: giovanni <giovanni.canali98@yahoo.it>

tests/test_callback/test_metric_tracker.py

@@ -23,17 +23,18 @@ def test_metric_tracker_constructor():
     MetricTracker()
 
 
-# def test_metric_tracker_routine(): #TODO revert
-#     # make the trainer
-#     trainer = Trainer(solver=solver,
-#                       callback=[
-#                           MetricTracker()
-#                       ],
-#                       accelerator='cpu',
-#                       max_epochs=5)
-#     trainer.train()
-#     # get the tracked metrics
-#     metrics = trainer.callback[0].metrics
-#     # assert the logged metrics are correct
-#     logged_metrics = sorted(list(metrics.keys()))
-#     assert logged_metrics == ['train_loss_epoch', 'train_loss_step', 'val_loss']
+def test_metric_tracker_routine():
+    # make the trainer
+    trainer = Trainer(
+        solver=solver,
+        callbacks=[MetricTracker()],
+        accelerator="cpu",
+        max_epochs=5,
+        log_every_n_steps=1,
+    )
+    trainer.train()
+    # get the tracked metrics
+    metrics = trainer.callbacks[0].metrics
+    # assert the logged metrics are correct
+    logged_metrics = sorted(list(metrics.keys()))
+    assert logged_metrics == ["train_loss"]

tests/test_callback/test_optimizer_callback.py

@@ -21,19 +21,25 @@ model = FeedForward(
 # make the solver
 solver = PINN(problem=poisson_problem, model=model)
 
-adam_optimizer = TorchOptimizer(torch.optim.Adam, lr=0.01)
-lbfgs_optimizer = TorchOptimizer(torch.optim.LBFGS, lr=0.001)
+adam = TorchOptimizer(torch.optim.Adam, lr=0.01)
+lbfgs = TorchOptimizer(torch.optim.LBFGS, lr=0.001)
 
 
 def test_switch_optimizer_constructor():
-    SwitchOptimizer(adam_optimizer, epoch_switch=10)
+    SwitchOptimizer(adam, epoch_switch=10)
 
 
-# def test_switch_optimizer_routine(): #TODO revert
-#     # make the trainer
-#     switch_opt_callback = SwitchOptimizer(lbfgs_optimizer, epoch_switch=3)
-#     trainer = Trainer(solver=solver,
-#                       callback=[switch_opt_callback],
-#                       accelerator='cpu',
-#                       max_epochs=5)
-#     trainer.train()
+def test_switch_optimizer_routine():
+    # check initial optimizer
+    solver.configure_optimizers()
+    assert solver.optimizer.instance.__class__ == torch.optim.Adam
+    # make the trainer
+    switch_opt_callback = SwitchOptimizer(lbfgs, epoch_switch=3)
+    trainer = Trainer(
+        solver=solver,
+        callbacks=[switch_opt_callback],
+        accelerator="cpu",
+        max_epochs=5,
+    )
+    trainer.train()
+    assert solver.optimizer.instance.__class__ == torch.optim.LBFGS

tests/test_callback/test_progress_bar.py

@@ -5,29 +5,32 @@ from pina.callback.processing_callback import PINAProgressBar
 from pina.problem.zoo import Poisson2DSquareProblem as Poisson
 
 
-# # make the problem
-# poisson_problem = Poisson()
-# boundaries = ['nil_g1', 'nil_g2', 'nil_g3', 'nil_g4']
-# n = 10
-# poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
-# poisson_problem.discretise_domain(n, 'grid', locations='laplace_D')
-# model = FeedForward(len(poisson_problem.input_variables),
-#                     len(poisson_problem.output_variables))
+# make the problem
+poisson_problem = Poisson()
+boundaries = ["g1", "g2", "g3", "g4"]
+n = 10
+condition_names = list(poisson_problem.conditions.keys())
+poisson_problem.discretise_domain(n, "grid", domains=boundaries)
+poisson_problem.discretise_domain(n, "grid", domains="D")
+model = FeedForward(
+    len(poisson_problem.input_variables), len(poisson_problem.output_variables)
+)
 
-# # make the solver
-# solver = PINN(problem=poisson_problem, model=model)
+# make the solver
+solver = PINN(problem=poisson_problem, model=model)
 
 
-# def test_progress_bar_constructor():
-#     PINAProgressBar(['mean'])
+def test_progress_bar_constructor():
+    PINAProgressBar()
 
-# def test_progress_bar_routine():
-#     # make the trainer
-#     trainer = Trainer(solver=solver,
-#                       callback=[
-#                           PINAProgressBar(['mean', 'laplace_D'])
-#                       ],
-#                       accelerator='cpu',
-#                       max_epochs=5)
-#     trainer.train()
-#     # TODO there should be a check that the correct metrics are displayed
+
+def test_progress_bar_routine():
+    # make the trainer
+    trainer = Trainer(
+        solver=solver,
+        callbacks=[PINAProgressBar(["val", condition_names[0]])],
+        accelerator="cpu",
+        max_epochs=5,
+    )
+    trainer.train()
+    # TODO there should be a check that the correct metrics are displayed