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: