158 lines
5.6 KiB
Python
158 lines
5.6 KiB
Python
import sys
|
|
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 problems.parametric_poisson import ParametricPoisson2DProblem as Poisson2D
|
|
from pina.deep_feed_forward import DeepFeedForward
|
|
|
|
from pina.adaptive_functions import AdaptiveSin, AdaptiveCos, AdaptiveTanh
|
|
|
|
|
|
class myFeature(torch.nn.Module):
|
|
"""
|
|
Feature: sin(x)
|
|
"""
|
|
|
|
def __init__(self):
|
|
super(myFeature, self).__init__()
|
|
|
|
def forward(self, x):
|
|
return torch.exp(- 2*(x[:, 0] - x[:, 2])**2 - 2*(x[:, 1] - x[:, 3])**2)
|
|
|
|
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 = Poisson2D()
|
|
model = DeepFeedForward(
|
|
layers=[200, 40, 10],
|
|
output_variables=poisson_problem.output_variables,
|
|
input_variables=poisson_problem.input_variables+['mu1', 'mu2'],
|
|
func=Softplus,
|
|
extra_features=feat
|
|
)
|
|
|
|
pinn = pPINN(
|
|
poisson_problem,
|
|
model,
|
|
lr=0.0006,
|
|
regularizer=1e-6,
|
|
lr_accelerate=None)
|
|
|
|
if args.s:
|
|
|
|
pinn.span_pts(30, 'chebyshev', ['D'])
|
|
pinn.span_pts(50, 'grid', ['gamma1', 'gamma2', 'gamma3', 'gamma4'])
|
|
#pinn.plot_pts()
|
|
pinn.train(10000, 10)
|
|
pinn.save_state('pina.poisson_param')
|
|
|
|
else:
|
|
pinn.load_state('pina.poisson_param')
|
|
#pinn.plot(40, torch.tensor([-0.8, -0.8]))
|
|
#pinn.plot(40, torch.tensor([ 0.8, 0.8]))
|
|
|
|
from smithers.io import VTUHandler
|
|
from scipy.interpolate import griddata
|
|
import matplotlib
|
|
matplotlib.use('GTK3Agg')
|
|
import matplotlib.pyplot as plt
|
|
|
|
res = 64
|
|
fname_minus = 'Poisson_param_08minus000000.vtu'
|
|
param = torch.tensor([-0.8, -0.8])
|
|
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]).reshape(-1, 2)], axis=1)
|
|
|
|
#unrolled_pts.to(dtype=self.dtype)
|
|
unrolled_pts = LabelTensor(unrolled_pts, ['x1', 'x2', 'mu1', 'mu2'])
|
|
|
|
Z_pred = pinn.model(unrolled_pts.tensor)
|
|
data = VTUHandler.read(fname_minus)
|
|
|
|
|
|
print(data['points'][:, :2].shape)
|
|
print(data['point_data']['f_16'].shape)
|
|
print(grids_container[0].shape)
|
|
print(grids_container[1].shape)
|
|
Z_truth = griddata(data['points'][:, :2], data['point_data']['f_16'], (grids_container[0], grids_container[1]))
|
|
|
|
|
|
err = np.abs(Z_truth + Z_pred.tensor.reshape(res, res).detach().numpy())
|
|
|
|
plt.subplot(1, 3, 1)
|
|
plt.pcolor(-Z_pred.tensor.reshape(res, res).detach())
|
|
plt.colorbar()
|
|
plt.subplot(1, 3, 2)
|
|
plt.pcolor(Z_truth)
|
|
plt.colorbar()
|
|
plt.subplot(1, 3, 3)
|
|
plt.pcolor(err, vmin=0, vmax=0.009)
|
|
plt.colorbar()
|
|
plt.show()
|
|
|
|
print(unrolled_pts.tensor.shape)
|
|
with open('parpoisson_minus_plot.txt', 'w') as f_:
|
|
f_.write('x y truth pred e\n')
|
|
for (x, y), tru, pre, e in zip(unrolled_pts[:, :2], Z_truth.reshape(-1, 1), -Z_pred.tensor.reshape(-1, 1), err.reshape(-1, 1)):
|
|
f_.write('{} {} {} {} {}\n'.format(x.item(), y.item(), tru.item(), pre.item(), e.item()))
|
|
|
|
fname_plus = 'Poisson_param_08plus000000.vtu'
|
|
param = torch.tensor([0.8, 0.8])
|
|
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]).reshape(-1, 2)], axis=1)
|
|
|
|
#unrolled_pts.to(dtype=self.dtype)
|
|
unrolled_pts = LabelTensor(unrolled_pts, ['x1', 'x2', 'mu1', 'mu2'])
|
|
|
|
Z_pred = pinn.model(unrolled_pts.tensor)
|
|
data = VTUHandler.read(fname_plus)
|
|
|
|
|
|
print(data['points'][:, :2].shape)
|
|
print(data['point_data']['f_16'].shape)
|
|
print(grids_container[0].shape)
|
|
print(grids_container[1].shape)
|
|
Z_truth = griddata(data['points'][:, :2], data['point_data']['f_16'], (grids_container[0], grids_container[1]))
|
|
|
|
|
|
err = np.abs(Z_truth + Z_pred.tensor.reshape(res, res).detach().numpy())
|
|
|
|
plt.subplot(1, 3, 1)
|
|
plt.pcolor(-Z_pred.tensor.reshape(res, res).detach())
|
|
plt.colorbar()
|
|
plt.subplot(1, 3, 2)
|
|
plt.pcolor(Z_truth)
|
|
plt.colorbar()
|
|
plt.subplot(1, 3, 3)
|
|
print('gggggg')
|
|
plt.pcolor(err, vmin=0, vmax=0.001)
|
|
plt.colorbar()
|
|
plt.show()
|
|
with open('parpoisson_plus_plot.txt', 'w') as f_:
|
|
f_.write('x y truth pred e\n')
|
|
for (x, y), tru, pre, e in zip(unrolled_pts[:, :2], Z_truth.reshape(-1, 1), -Z_pred.tensor.reshape(-1, 1), err.reshape(-1, 1)):
|
|
f_.write('{} {} {} {} {}\n'.format(x.item(), y.item(), tru.item(), pre.item(), e.item()))
|
|
|
|
|