43 lines
1.3 KiB
Python
43 lines
1.3 KiB
Python
import argparse
|
|
|
|
from torch.nn import Softplus
|
|
|
|
from pina.model import FeedForward
|
|
from pina import Plotter, PINN
|
|
from problems.first_order_ode import FirstOrderODE
|
|
|
|
|
|
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)
|
|
args = parser.parse_args()
|
|
|
|
# define Problem + Model + PINN
|
|
problem = FirstOrderODE()
|
|
model = FeedForward(
|
|
layers=[4]*2,
|
|
output_variables=problem.output_variables,
|
|
input_variables=problem.input_variables,
|
|
func=Softplus,
|
|
)
|
|
pinn = PINN(problem, model, lr=0.03, error_norm='mse', regularizer=0)
|
|
|
|
if args.s:
|
|
|
|
pinn.span_pts(
|
|
{'variables': ['x'], 'mode': 'grid', 'n': 1}, locations=['bc'])
|
|
pinn.span_pts(
|
|
{'variables': ['x'], 'mode': 'grid', 'n': 30}, locations=['dd'])
|
|
Plotter().plot_samples(pinn, ['x'])
|
|
pinn.train(1200, 50)
|
|
pinn.save_state('pina.ode')
|
|
|
|
else:
|
|
pinn.load_state('pina.ode')
|
|
plotter = Plotter()
|
|
plotter.plot(pinn, components=['y'])
|