diff --git a/tests/test_solvers/test_supervised_solver.py b/tests/test_solvers/test_supervised_solver.py new file mode 100644 index 0000000..dfe0bd8 --- /dev/null +++ b/tests/test_solvers/test_supervised_solver.py @@ -0,0 +1,101 @@ +import torch + +from pina.problem import AbstractProblem +from pina import Condition, LabelTensor +from pina.solvers import SupervisedSolver +from pina.trainer import Trainer +from pina.model import FeedForward +from pina.loss import LpLoss + + +class NeuralOperatorProblem(AbstractProblem): + input_variables = ['u_0', 'u_1'] + output_variables = ['u'] + conditions = {'data' : Condition(input_points=LabelTensor(torch.rand(100, 2), input_variables), + output_points=LabelTensor(torch.rand(100, 1), output_variables))} + +class myFeature(torch.nn.Module): + """ + Feature: sin(x) + """ + + def __init__(self): + super(myFeature, self).__init__() + + def forward(self, x): + t = (torch.sin(x.extract(['u_0']) * torch.pi) * + torch.sin(x.extract(['u_1']) * torch.pi)) + return LabelTensor(t, ['sin(x)sin(y)']) + + +# make the problem + extra feats +problem = NeuralOperatorProblem() +extra_feats = [myFeature()] +model = FeedForward(len(problem.input_variables), + len(problem.output_variables)) +model_extra_feats = FeedForward( + len(problem.input_variables) + 1, + len(problem.output_variables)) + + +def test_constructor(): + SupervisedSolver(problem=problem, model=model, extra_features=None) + + +def test_constructor_extra_feats(): + SupervisedSolver(problem=problem, model=model_extra_feats, extra_features=extra_feats) + + +def test_train_cpu(): + solver = SupervisedSolver(problem = problem, model=model, extra_features=None, loss=LpLoss()) + trainer = Trainer(solver=solver, max_epochs=3, accelerator='cpu', batch_size=20) + trainer.train() + + +def test_train_restore(): + tmpdir = "tests/tmp_restore" + solver = SupervisedSolver(problem=problem, + model=model, + extra_features=None, + loss=LpLoss()) + trainer = Trainer(solver=solver, + max_epochs=5, + accelerator='cpu', + default_root_dir=tmpdir) + trainer.train() + ntrainer = Trainer(solver=solver, 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" + solver = SupervisedSolver(problem=problem, + model=model, + extra_features=None, + loss=LpLoss()) + trainer = Trainer(solver=solver, + max_epochs=15, + accelerator='cpu', + default_root_dir=tmpdir) + trainer.train() + new_solver = SupervisedSolver.load_from_checkpoint( + f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=15.ckpt', + problem = problem, model=model) + test_pts = LabelTensor(torch.rand(20, 2), problem.input_variables) + assert new_solver.forward(test_pts).shape == (20, 1) + assert new_solver.forward(test_pts).shape == solver.forward(test_pts).shape + torch.testing.assert_close( + new_solver.forward(test_pts), + solver.forward(test_pts)) + import shutil + shutil.rmtree(tmpdir) + +def test_train_extra_feats_cpu(): + pinn = SupervisedSolver(problem=problem, + model=model_extra_feats, + extra_features=extra_feats) + trainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu') + trainer.train() \ No newline at end of file