def test(): import pdedata import torchvision trans = torchvision.transforms.Compose( [pdedata.DownSample(4), pdedata.ToTensor()]) d = pdedata.variantcoelinear2d(0.6, mesh_size=[200, 200], initfreq=4, transform=trans) dataloader = torch.utils.data.DataLoader(d, batch_size=2, num_workers=2) dataloader = iter(dataloader) sample = next(dataloader) sample = pdedata.ToVariable()(sample) xy = torch.stack([sample['x'], sample['y']], dim=3) u0 = sample['u0'] linpdelearner = VariantCoeLinear2d(kernel_size=[7, 7], max_order=4, dx=2 * pi / 50, constraint='moment', xy=xy, interp_degree=2, interp_mesh_size=[20, 20], dt=1e-2) linpdelearner.xy = xy.clone() ut = linpdelearner(u0, 20) trans = torchvision.transforms.Compose( [pdedata.DownSample(4, boundary='Dirichlet'), pdedata.ToTensor()]) d = pdedata.singlenonlinear2d(0.6, mesh_size=[200, 200], transform=trans) dataloader = torch.utils.data.DataLoader(d, batch_size=2, num_workers=2) dataloader = iter(dataloader) sample = next(dataloader) sample = pdedata.ToVariable()(sample) xy = torch.stack([sample['x'], sample['y']], dim=3) u0 = sample['u0'] nonlinearlearner = SingleNonLinear2d(kernel_size=[7, 7], max_order=2, dx=2 * pi / 50, constraint='moment', xy=xy, interp_degree=2, interp_mesh_size=[5, 5], nonlinear_interp_degree=3, nonlinear_interp_mesh_bound=[-30, 30], nonlinear_interp_mesh_size=40, dt=1e-2) nonlinearlearner.xy = xy.clone() ut = nonlinearlearner(u0, 20)
j = 0 nonlintermtest = nonlinpdelearner.xy.data.new(2000) nonlintermtest.copy_( torch.linspace(nonlinear_interp_mesh_bound[0], nonlinear_interp_mesh_bound[1], 2000)) testresults = nonlinpdelearner.nonlinearfitter(Variable(nonlintermtest)) testresults = testresults.data.cpu().numpy() nonlintermtest = nonlintermtest.cpu().numpy() a = pltnewaxis() a.plot(nonlintermtest, testresults) a.plot(nonlintermtest, np.sin(nonlintermtest) * nonlinear_coefficient) #%% show prediction trans = torchvision.transforms.Compose([ pdedata.DownSample(5, 'Dirichlet'), pdedata.ToTensor(), pdedata.ToPrecision(precision), pdedata.AddNoise(start_noise_level, end_noise_level) ]) d = pdedata.singlenonlinear2d(np.array(range(1, teststepnum + 1)) * dt, mesh_size=mesh_size * 5, initfreq=initfreq, diffusivity=diffusivity, nonlinear_coefficient=nonlinear_coefficient, transform=trans) sample = pdedata.ToVariable()(d[0]) xy = torch.stack([sample['x'], sample['y']], dim=2) u0 = sample['u0'] uT = sample['uT'] u0 = (u0.cuda(gpu) if gpu >= 0 else u0) nonlinpdelearner.id.MomentBank.moment.volatile = True
coe = eval('linpdelearner.coe'+str(i)) coe = coe().data.cpu().numpy() ax = a.flatten()[i] b = ax.imshow(coe, cmap='jet', vmin=-5, vmax=5) ax.set_xticks([]) ax.set_yticks([]) ax.set_title('C'+str(k-j)+str(j), fontsize=20) c = ax.figure.colorbar(b,ax=ax) c.set_ticks([-5,-2.5,0,2.5,5]) j += 1 if j > k: k += 1 j = 0 #%% show prediction trans = torchvision.transforms.Compose([pdedata.DownSample(5), pdedata.ToTensor(), pdedata.ToPrecision(precision), pdedata.AddNoise(start_noise_level, end_noise_level)]) d = pdedata.variantcoelinear2d(np.array(range(1,teststepnum+1))*dt, mesh_size=mesh_size*5, initfreq=initfreq, variant_coe_magnitude=variant_coe_magnitude, transform=trans) sample = pdedata.ToVariable()(d[0]) xy = torch.stack([sample['x'],sample['y']], dim=2) u0 = sample['u0'] uT = sample['uT'] u0 = (u0.cuda(gpu) if gpu>=0 else u0) linpdelearner.id.MomentBank.moment.volatile = True linpdelearner.xy = xy # set xy for pde-net F = pltnewmeshbar((3,5)) i = 0 for j in (0,20,40,60,80): if j == 0: u_true = u0.data.cpu().numpy() u = u_true else: