From 1b14b55f14534e2429e8abfce00f76560ee4d4d7 Mon Sep 17 00:00:00 2001
From: Nicola Demo <demo.nicola@gmail.com>
Date: Tue, 25 Jul 2023 17:23:12 +0200
Subject: [PATCH] Trainer `train` simplified, tests for load  (#168)

- the arguments of Trainer.train now are passed to the fit
- unittest for load/restoring from checkpoint
---
 pina/trainer.py                  |  6 ++---
 tests/test_solvers/test_garom.py |  2 +-
 tests/test_solvers/test_pinn.py  | 43 +++++++++++++++++++++++++++++---
 3 files changed, 44 insertions(+), 7 deletions(-)

diff --git a/pina/trainer.py b/pina/trainer.py
index 77610a5d..44c4b11e 100644
--- a/pina/trainer.py
+++ b/pina/trainer.py
@@ -7,7 +7,7 @@
 
 class Trainer(pl.Trainer):
 
-    def __init__(self, solver, kwargs={}):
+    def __init__(self, solver, **kwargs):
         super().__init__(**kwargs)
         
         # get accellerator
@@ -29,6 +29,6 @@ def __init__(self, solver, kwargs={}):
         self._loader = DummyLoader(solver.problem.input_pts, device) 
 
 
-    def train(self): # TODO add kwargs and lightining capabilities
-        return super().fit(self._model, self._loader)
+    def train(self, **kwargs): # TODO add kwargs and lightining capabilities
+        return super().fit(self._model, self._loader, **kwargs)
     
diff --git a/tests/test_solvers/test_garom.py b/tests/test_solvers/test_garom.py
index d83bfc1f..754dfd76 100644
--- a/tests/test_solvers/test_garom.py
+++ b/tests/test_solvers/test_garom.py
@@ -134,7 +134,7 @@ def test_train_cpu():
                                                 hidden_dimension=64)
                 )
 
-    trainer = Trainer(solver=solver, kwargs={'max_epochs' : 4, 'accelerator': 'cpu'})
+    trainer = Trainer(solver=solver, max_epochs=4, accelerator='cpu')
     trainer.train()
 
 def test_sample():
diff --git a/tests/test_solvers/test_pinn.py b/tests/test_solvers/test_pinn.py
index 92ce0b37..c8dcee26 100644
--- a/tests/test_solvers/test_pinn.py
+++ b/tests/test_solvers/test_pinn.py
@@ -56,12 +56,12 @@ def poisson_sol(self, pts):
 
     truth_solution = poisson_sol
 
-
 class myFeature(torch.nn.Module):
     """
     Feature: sin(x)
     """
 
+
     def __init__(self):
         super(myFeature, self).__init__()
 
@@ -92,8 +92,45 @@ def test_train_cpu():
     n = 10
     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
     pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
-    trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'cpu'})
+    trainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu')
+    trainer.train()
+
+def test_train_restore():
+    tmpdir = "tests/tmp_restore"
+    poisson_problem = Poisson()
+    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+    n = 10
+    poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
+    pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
+    trainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu', default_root_dir=tmpdir)
+    trainer.train()
+    print('ggg')
+    ntrainer = Trainer(solver=pinn, max_epochs=15, accelerator='cpu')
+    t = ntrainer.train(
+        ckpt_path=f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=5.ckpt')
+    import shutil
+    shutil.rmtree(tmpdir)
+
+def test_train_load():
+    tmpdir = "tests/tmp_load"
+    poisson_problem = Poisson()
+    boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
+    n = 10
+    poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
+    pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss())
+    trainer = Trainer(solver=pinn, max_epochs=15, accelerator='cpu',
+                     default_root_dir=tmpdir)
     trainer.train()
+    new_pinn = PINN.load_from_checkpoint(
+        f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=15.ckpt',
+        problem = poisson_problem, model=model)
+    test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10)
+    assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1)
+    assert new_pinn.forward(test_pts).extract(['u']).shape == pinn.forward(test_pts).extract(['u']).shape
+    torch.testing.assert_close(new_pinn.forward(test_pts).extract(['u']), pinn.forward(test_pts).extract(['u']))
+    import shutil
+    shutil.rmtree(tmpdir)
+    
 
 # # TODO fix asap. Basically sampling few variables
 # # works only if both variables are in a range.
@@ -118,7 +155,7 @@ def test_train_extra_feats_cpu():
     n = 10
     poisson_problem.discretise_domain(n, 'grid', locations=boundaries)
     pinn = PINN(problem = poisson_problem, model=model_extra_feats, extra_features=extra_feats)
-    trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'cpu'})
+    trainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu')
     trainer.train()
 
 # TODO, fix GitHub actions to run also on GPU