90 lines
2.9 KiB
Python
90 lines
2.9 KiB
Python
import argparse
|
|
import numpy as np
|
|
import torch
|
|
from torch.nn import Softplus
|
|
|
|
from pina import LabelTensor
|
|
from pina.solvers import PINN
|
|
from pina.model import MultiFeedForward, FeedForward
|
|
from pina.plotter import Plotter
|
|
from pina.trainer import Trainer
|
|
from problems.parametric_elliptic_optimal_control import (
|
|
ParametricEllipticOptimalControl)
|
|
|
|
|
|
class myFeature(torch.nn.Module):
|
|
"""
|
|
Feature: sin(x)
|
|
"""
|
|
|
|
def __init__(self):
|
|
super(myFeature, self).__init__()
|
|
|
|
def forward(self, x):
|
|
t = (-x.extract(['x1'])**2+1) * (-x.extract(['x2'])**2+1)
|
|
return LabelTensor(t, ['k0'])
|
|
|
|
|
|
class PIArch(MultiFeedForward):
|
|
|
|
def __init__(self, dff_dict):
|
|
super().__init__(dff_dict)
|
|
|
|
def forward(self, x):
|
|
out = self.uy(x)
|
|
out.labels = ['u', 'y']
|
|
z = LabelTensor(
|
|
(out.extract(['u']) * x.extract(['alpha'])), ['z'])
|
|
return out.append(z)
|
|
|
|
if __name__ == "__main__":
|
|
|
|
parser = argparse.ArgumentParser(description="Run PINA")
|
|
parser.add_argument("--load", help="directory to save or load file", type=str)
|
|
parser.add_argument("--features", help="extra features", type=int)
|
|
parser.add_argument("--epochs", help="extra features", type=int, default=1000)
|
|
args = parser.parse_args()
|
|
|
|
if args.features is None:
|
|
args.features = 0
|
|
|
|
# extra features
|
|
feat = [myFeature()] if args.features else []
|
|
args = parser.parse_args()
|
|
|
|
# create problem and discretise domain
|
|
opc = ParametricEllipticOptimalControl()
|
|
opc.discretise_domain(n= 900, mode='random', variables=['x1', 'x2'], locations=['D'])
|
|
opc.discretise_domain(n= 5, mode='random', variables=['mu', 'alpha'], locations=['D'])
|
|
opc.discretise_domain(n= 200, mode='random', variables=['x1', 'x2'], locations=['gamma1', 'gamma2', 'gamma3', 'gamma4'])
|
|
opc.discretise_domain(n= 5, mode='random', variables=['mu', 'alpha'], locations=['gamma1', 'gamma2', 'gamma3', 'gamma4'])
|
|
|
|
# create model
|
|
model = PIArch(
|
|
{
|
|
'uy': {
|
|
'input_dimensions': 4 + len(feat),
|
|
'output_dimensions': 2,
|
|
'layers': [40, 40, 20],
|
|
'func': Softplus,
|
|
},
|
|
}
|
|
)
|
|
|
|
# create PINN
|
|
pinn = PINN(problem=opc, model=model, optimizer_kwargs={'lr' : 0.002}, extra_features=feat)
|
|
|
|
# create trainer
|
|
directory = 'pina.parametric_optimal_control_{}'.format(bool(args.features))
|
|
trainer = Trainer(solver=pinn, accelerator='cpu', max_epochs=args.epochs, default_root_dir=directory)
|
|
|
|
|
|
if args.load:
|
|
pinn = PINN.load_from_checkpoint(checkpoint_path=args.load, problem=opc, model=model, extra_features=feat)
|
|
plotter = Plotter()
|
|
plotter.plot(pinn, fixed_variables={'mu' : 3 , 'alpha' : 1}, components='u')
|
|
plotter.plot(pinn, fixed_variables={'mu' : 3 , 'alpha' : 1}, components='z')
|
|
plotter.plot(pinn, fixed_variables={'mu' : 3 , 'alpha' : 1}, components='y')
|
|
else:
|
|
trainer.train()
|