40 lines
1.2 KiB
Python
40 lines
1.2 KiB
Python
from pina.solver import PINN
|
|
from pina.trainer import Trainer
|
|
from pina.model import FeedForward
|
|
from pina.callback import MetricTracker
|
|
from pina.problem.zoo import Poisson2DSquareProblem as Poisson
|
|
|
|
|
|
# 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)
|
|
|
|
|
|
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']
|