66 lines
2.1 KiB
Python
66 lines
2.1 KiB
Python
import argparse
|
|
import torch
|
|
from torch.nn import Softplus
|
|
from pina import Plotter, LabelTensor, PINN
|
|
from pina.model import FeedForward
|
|
from problems.parametric_poisson import ParametricPoisson
|
|
|
|
|
|
class myFeature(torch.nn.Module):
|
|
"""
|
|
"""
|
|
def __init__(self):
|
|
super(myFeature, self).__init__()
|
|
|
|
def forward(self, x):
|
|
t = (
|
|
torch.exp(
|
|
- 2*(x.extract(['x']) - x.extract(['mu1']))**2
|
|
- 2*(x.extract(['y']) - x.extract(['mu2']))**2
|
|
)
|
|
)
|
|
return LabelTensor(t, ['k0'])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
parser = argparse.ArgumentParser(description="Run PINA")
|
|
group = parser.add_mutually_exclusive_group(required=True)
|
|
group.add_argument("-s", "-save", action="store_true")
|
|
group.add_argument("-l", "-load", action="store_true")
|
|
parser.add_argument("id_run", help="number of run", type=int)
|
|
parser.add_argument("features", help="extra features", type=int)
|
|
args = parser.parse_args()
|
|
|
|
feat = [myFeature()] if args.features else []
|
|
|
|
poisson_problem = ParametricPoisson()
|
|
model = FeedForward(
|
|
layers=[10, 10, 10],
|
|
output_variables=poisson_problem.output_variables,
|
|
input_variables=poisson_problem.input_variables,
|
|
func=Softplus,
|
|
extra_features=feat
|
|
)
|
|
|
|
pinn = PINN(poisson_problem, model, lr=0.006, regularizer=1e-6)
|
|
|
|
if args.s:
|
|
|
|
pinn.span_pts(
|
|
{'variables': ['x', 'y'], 'mode': 'random', 'n': 100},
|
|
{'variables': ['mu1', 'mu2'], 'mode': 'grid', 'n': 5},
|
|
locations=['D'])
|
|
pinn.span_pts(
|
|
{'variables': ['x', 'y'], 'mode': 'grid', 'n': 20},
|
|
{'variables': ['mu1', 'mu2'], 'mode': 'grid', 'n': 5},
|
|
locations=['gamma1', 'gamma2', 'gamma3', 'gamma4'])
|
|
pinn.train(10000, 100)
|
|
pinn.save_state('pina.poisson_param')
|
|
|
|
else:
|
|
pinn.load_state('pina.poisson_param')
|
|
plotter = Plotter()
|
|
plotter.plot(pinn, fixed_variables={'mu1': 0, 'mu2': 1}, levels=21)
|
|
plotter.plot(pinn, fixed_variables={'mu1': 1, 'mu2': -1}, levels=21)
|