diff --git a/examples/ldc/ldc.py b/examples/ldc/ldc.py new file mode 100755 index 0000000..1cd8a31 --- /dev/null +++ b/examples/ldc/ldc.py @@ -0,0 +1,115 @@ +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()