import numpy as np import torch import argparse from pina.pinn import PINN from pina.ppinn import ParametricPINN as pPINN from pina.label_tensor import LabelTensor from torch.nn import ReLU, Tanh, Softplus from pina.adaptive_functions.adaptive_softplus import AdaptiveSoftplus from problems.parametric_elliptic_optimal_control_alpha_variable import ParametricEllipticOptimalControl from pina.multi_deep_feed_forward import MultiDeepFeedForward from pina.deep_feed_forward import DeepFeedForward alpha = 1 class myFeature(torch.nn.Module): """ Feature: sin(x) """ def __init__(self): super(myFeature, self).__init__() def forward(self, x): return (-x[:, 0]**2+1) * (-x[:, 1]**2+1) class CustomMultiDFF(MultiDeepFeedForward): def __init__(self, dff_dict): super().__init__(dff_dict) def forward(self, x): out = self.uu(x) p = LabelTensor((out['u_param'] * x[:, 3]).reshape(-1, 1), ['p']) a = LabelTensor.hstack([out, p]) return a model = CustomMultiDFF( { 'uu': { 'input_variables': ['x1', 'x2', 'mu', 'alpha'], 'output_variables': ['u_param', 'y'], 'layers': [40, 40, 20], 'func': Softplus, 'extra_features': [myFeature()], }, # 'u_param': { # 'input_variables': ['u', 'mu'], # 'output_variables': ['u_param'], # 'layers': [], # 'func': None # }, # 'p': { # 'input_variables': ['u'], # 'output_variables': ['p'], # 'layers': [10], # 'func': None # }, } ) opc = ParametricEllipticOptimalControl(alpha) 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") args = parser.parse_args() # model = DeepFeedForward( # layers=[40, 40, 20], # output_variables=['u_param', 'y', 'p'], # input_variables=opc.input_variables+['mu', 'alpha'], # func=Softplus, # extra_features=[myFeature()] # ) pinn = pPINN( opc, model, lr=0.002, error_norm='mse', regularizer=1e-8, lr_accelerate=None) if args.s: pinn.span_pts(30, 'grid', ['D1']) pinn.span_pts(50, 'grid', ['gamma1', 'gamma2', 'gamma3', 'gamma4']) pinn.train(10000, 20) # with open('ocp_wrong_history.txt', 'w') as file_: # for i, losses in enumerate(pinn.history): # file_.write('{} {}\n'.format(i, sum(losses).item())) pinn.save_state('pina.ocp') else: pinn.load_state('working.pina.ocp') pinn.load_state('pina.ocp') import matplotlib matplotlib.use('GTK3Agg') import matplotlib.pyplot as plt # res = 64 # param = torch.tensor([[3., 1]]) # pts_container = [] # for mn, mx in [[-1, 1], [-1, 1]]: # pts_container.append(np.linspace(mn, mx, res)) # grids_container = np.meshgrid(*pts_container) # unrolled_pts = torch.tensor([t.flatten() for t in grids_container]).T # unrolled_pts = torch.cat([unrolled_pts, param.double().repeat(unrolled_pts.shape[0], 1).reshape(-1, 2)], axis=1) # unrolled_pts = LabelTensor(unrolled_pts, ['x1', 'x2', 'mu', 'alpha']) # Z_pred = pinn.model(unrolled_pts.tensor) # print(Z_pred.tensor.shape) # plt.subplot(2, 3, 1) # plt.pcolor(Z_pred['y'].reshape(res, res).detach()) # plt.colorbar() # plt.subplot(2, 3, 2) # plt.pcolor(Z_pred['u_param'].reshape(res, res).detach()) # plt.colorbar() # plt.subplot(2, 3, 3) # plt.pcolor(Z_pred['p'].reshape(res, res).detach()) # plt.colorbar() # with open('ocp_mu3_a1_plot.txt', 'w') as f_: # f_.write('x y u p ys\n') # for (x, y), tru, pre, e in zip(unrolled_pts[:, :2], # Z_pred['u_param'].reshape(-1, 1), # Z_pred['p'].reshape(-1, 1), # Z_pred['y'].reshape(-1, 1), # ): # f_.write('{} {} {} {} {}\n'.format(x.item(), y.item(), tru.item(), pre.item(), e.item())) # param = torch.tensor([[3.0, 0.01]]) # unrolled_pts = torch.tensor([t.flatten() for t in grids_container]).T # unrolled_pts = torch.cat([unrolled_pts, param.double().repeat(unrolled_pts.shape[0], 1).reshape(-1, 2)], axis=1) # unrolled_pts = LabelTensor(unrolled_pts, ['x1', 'x2', 'mu', 'alpha']) # Z_pred = pinn.model(unrolled_pts.tensor) # plt.subplot(2, 3, 4) # plt.pcolor(Z_pred['y'].reshape(res, res).detach()) # plt.colorbar() # plt.subplot(2, 3, 5) # plt.pcolor(Z_pred['u_param'].reshape(res, res).detach()) # plt.colorbar() # plt.subplot(2, 3, 6) # plt.pcolor(Z_pred['p'].reshape(res, res).detach()) # plt.colorbar() # plt.show() # with open('ocp_mu3_a0.01_plot.txt', 'w') as f_: # f_.write('x y u p ys\n') # for (x, y), tru, pre, e in zip(unrolled_pts[:, :2], # Z_pred['u_param'].reshape(-1, 1), # Z_pred['p'].reshape(-1, 1), # Z_pred['y'].reshape(-1, 1), # ): # f_.write('{} {} {} {} {}\n'.format(x.item(), y.item(), tru.item(), pre.item(), e.item())) y = {} u = {} for alpha in [0.01, 0.1, 1]: y[alpha] = [] u[alpha] = [] for p in np.linspace(0.5, 3, 32): a = pinn.model(LabelTensor(torch.tensor([[0, 0, p, alpha]]).double(), ['x1', 'x2', 'mu', 'alpha']).tensor) y[alpha].append(a['y'].detach().numpy()[0]) u[alpha].append(a['u_param'].detach().numpy()[0]) plt.plot(np.linspace(0.5, 3, 32), u[1], label='u') plt.plot(np.linspace(0.5, 3, 32), u[0.01], label='u') plt.plot(np.linspace(0.5, 3, 32), u[0.1], label='u') plt.plot([1, 2, 3], [0.28, 0.56, 0.85], 'o', label='Truth values') plt.legend() plt.show() print(y[1]) print(y[0.1]) print(y[0.01]) with open('elliptic_param_y.txt', 'w') as f_: f_.write('mu 1 01 001\n') for mu, y1, y01, y001 in zip(np.linspace(0.5, 3, 32), y[1], y[0.1], y[0.01]): f_.write('{} {} {} {}\n'.format(mu, y1, y01, y001)) with open('elliptic_param_u.txt', 'w') as f_: f_.write('mu 1 01 001\n') for mu, y1, y01, y001 in zip(np.linspace(0.5, 3, 32), u[1], u[0.1], u[0.01]): f_.write('{} {} {} {}\n'.format(mu, y1, y01, y001)) plt.plot(np.linspace(0.5, 3, 32), y, label='y') plt.plot([1, 2, 3], [0.062, 0.12, 0.19], 'o', label='Truth values') plt.legend() plt.show()