46 lines
1.3 KiB
Python
46 lines
1.3 KiB
Python
from pina.callback import SwitchOptimizer
|
|
import torch
|
|
import pytest
|
|
|
|
from pina.solver import PINN
|
|
from pina.trainer import Trainer
|
|
from pina.model import FeedForward
|
|
from pina.problem.zoo import Poisson2DSquareProblem as Poisson
|
|
from pina.optim import TorchOptimizer
|
|
|
|
# make the problem
|
|
poisson_problem = Poisson()
|
|
boundaries = ["g1", "g2", "g3", "g4"]
|
|
n = 10
|
|
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)
|
|
|
|
adam = TorchOptimizer(torch.optim.Adam, lr=0.01)
|
|
lbfgs = TorchOptimizer(torch.optim.LBFGS, lr=0.001)
|
|
|
|
|
|
def test_switch_optimizer_constructor():
|
|
SwitchOptimizer(adam, epoch_switch=10)
|
|
|
|
|
|
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
|