forked from idrl/idrlnet
116 lines
4.3 KiB
Python
Executable File
116 lines
4.3 KiB
Python
Executable File
from idrlnet import shortcut as sc
|
|
import sympy as sp
|
|
import torch
|
|
import numpy as np
|
|
from idrlnet.pde_op.equations import NavierStokesNode
|
|
from matplotlib import tri
|
|
import matplotlib.pyplot as plt
|
|
|
|
x, y = sp.symbols('x y')
|
|
rec = sc.Rectangle((-0.05, -0.05), (0.05, 0.05))
|
|
|
|
|
|
@sc.datanode(name='flow_domain')
|
|
class InteriorDomain(sc.SampleDomain):
|
|
def __init__(self):
|
|
points = rec.sample_interior(400000, bounds={x: (-0.05, 0.05), y: (-0.05, 0.05)})
|
|
|
|
constraints = sc.Variables({'continuity': torch.zeros(len(points['x']), 1),
|
|
'momentum_x': torch.zeros(len(points['x']), 1),
|
|
'momentum_y': torch.zeros(len(points['x']), 1)})
|
|
|
|
points['area'] = np.ones_like(points['area']) * 2.5e-6
|
|
|
|
constraints['lambda_continuity'] = points['sdf']
|
|
constraints['lambda_momentum_x'] = points['sdf']
|
|
constraints['lambda_momentum_y'] = points['sdf']
|
|
self.points = sc.Variables(points).to_torch_tensor_()
|
|
self.constraints = sc.Variables(constraints).to_torch_tensor_()
|
|
|
|
def sampling(self, *args, **kwargs):
|
|
return self.points, self.constraints
|
|
|
|
|
|
@sc.datanode(name='left_right_down')
|
|
class LeftRightDownBoundaryDomain(sc.SampleDomain):
|
|
def __init__(self):
|
|
points = rec.sample_boundary(3333, sieve=(y < 0.05))
|
|
constraints = sc.Variables({'u': torch.zeros(len(points['x']), 1), 'v': torch.zeros(len(points['x']), 1)})
|
|
points['area'] = np.ones_like(points['area']) * 1e-4
|
|
|
|
self.points = sc.Variables(points).to_torch_tensor_()
|
|
self.constraints = sc.Variables(constraints).to_torch_tensor_()
|
|
|
|
def sampling(self, *args, **kwargs):
|
|
return self.points, self.constraints
|
|
|
|
|
|
@sc.datanode(name='up')
|
|
class UpBoundaryDomain(sc.SampleDomain):
|
|
def __init__(self):
|
|
points = rec.sample_boundary(10000, sieve=sp.Eq(y, 0.05))
|
|
points['area'] = np.ones_like(points['area']) * 1e-4
|
|
constraints = sc.Variables({'u': torch.ones(len(points['x']), 1), 'v': torch.zeros(len(points['x']), 1)})
|
|
constraints['lambda_u'] = 1 - 20 * abs(points['x'].copy())
|
|
self.points = sc.Variables(points).to_torch_tensor_()
|
|
self.constraints = sc.Variables(constraints).to_torch_tensor_()
|
|
|
|
def sampling(self, *args, **kwargs):
|
|
return self.points, self.constraints
|
|
|
|
|
|
torch.autograd.set_detect_anomaly(True)
|
|
net = sc.MLP([2, 100, 100, 100, 100, 3], activation=sc.Activation.tanh, initialization=sc.Initializer.Xavier_uniform,
|
|
weight_norm=False)
|
|
|
|
net_u = sc.NetNode(inputs=('x', 'y',), outputs=('u', 'v', 'p'), net=net, name='net_u')
|
|
pde = NavierStokesNode(nu=0.01, rho=1.0, dim=2, time=False)
|
|
|
|
s = sc.Solver(sample_domains=(InteriorDomain(), LeftRightDownBoundaryDomain(), UpBoundaryDomain()),
|
|
netnodes=[net_u],
|
|
pdes=[pde],
|
|
max_iter=4000,
|
|
network_dir='./result/tanh_Xavier_uniform',
|
|
)
|
|
s.solve()
|
|
|
|
|
|
def interoir_domain_infer():
|
|
points = rec.sample_interior(1000000, bounds={x: (-0.05, 0.05), y: (-0.05, 0.05)})
|
|
constraints = sc.Variables({'continuity': torch.zeros(len(points['x']), 1),
|
|
'momentum_x': torch.zeros(len(points['x']), 1),
|
|
'momentum_y': torch.zeros(len(points['x']), 1)})
|
|
return points, constraints
|
|
|
|
|
|
data_infer = sc.get_data_node(interoir_domain_infer, name='flow_domain')
|
|
s.sample_domains = [data_infer]
|
|
|
|
pred = s.infer_step({'flow_domain': ['x', 'y', 'v', 'u', 'p']})
|
|
num_x = pred['flow_domain']['x'].detach().cpu().numpy().ravel()
|
|
num_y = pred['flow_domain']['y'].detach().cpu().numpy().ravel()
|
|
num_u = pred['flow_domain']['u'].detach().cpu().numpy().ravel()
|
|
num_v = pred['flow_domain']['v'].detach().cpu().numpy().ravel()
|
|
num_p = pred['flow_domain']['p'].detach().cpu().numpy().ravel()
|
|
triang_total = tri.Triangulation(num_x, num_y)
|
|
|
|
triang_total = tri.Triangulation(num_x, num_y)
|
|
u_pre = num_u.flatten()
|
|
|
|
fig = plt.figure(figsize=(15, 5))
|
|
ax1 = fig.add_subplot(131)
|
|
tcf = ax1.tricontourf(triang_total, num_u, 100, cmap="jet")
|
|
tc_bar = plt.colorbar(tcf)
|
|
ax1.set_title('u')
|
|
|
|
ax2 = fig.add_subplot(132)
|
|
tcf = ax2.tricontourf(triang_total, num_v, 100, cmap="jet")
|
|
tc_bar = plt.colorbar(tcf)
|
|
ax2.set_title('v')
|
|
|
|
ax3 = fig.add_subplot(133)
|
|
tcf = ax3.tricontourf(triang_total, num_p, 100, cmap="jet")
|
|
tc_bar = plt.colorbar(tcf)
|
|
ax3.set_title('p')
|
|
plt.show()
|