def potential(x):
k = 4
return k*x**2/2
def compose_psi(x, N):
f_b = 0
dt = x - x_min
psi = f_b + (1 - torch.exp(-dt)) * (1 - torch.exp(dt-x_max*2)) * N
return psi
def driver(index):
return -2+0.25*(index//2000)
grid = torch.linspace(x_min, x_max, 400).reshape(-1, 1)
model = EigenvalueProblemModel([1, 50, 50, 1], Sin, compose_psi, PDE_loss, lr=8e-3, start_eigenvalue=1.0)
model.train(driver, 1000, grid, partial(perturb1D, x_min=x_min, x_max=x_max), int(4e5), max_required_loss=1e-2, rtol=0.001, fraction=3)
model.plot_history()
large_grid = torch.linspace(x_min, x_max, 400).reshape(-1, 1)
for marker in model.eigenfunctions.keys():
y = model.get_eigenfunction(marker)(large_grid)
plt.plot(large_grid.reshape(-1).numpy(), y.reshape(-1).detach().numpy())
plt.xlabel("x")
plt.ylabel("$\Psi(x)$")
plt.xlim(x_min, x_max)
plt.grid()
plt.title(f"Eigenfunction {marker} with energy {model.eigenfunctions[marker][1]}")
plt.show()
def driver(index):
return -3.5e-18 + 0.1e-18 * (index // 1000)
grid1D = torch.linspace(-bound, bound, 5)
grid_x, grid_y, grid_z = torch.meshgrid(grid1D, grid1D, grid1D)
grid = torch.cat(
[grid_x.reshape(-1, 1),
grid_y.reshape(-1, 1),
grid_z.reshape(-1, 1)],
dim=1)
model = EigenvalueProblemModel([3, 20, 20, 20, 20, 1],
Sin,
compose_psi,
PDE_loss,
lr=1e-20,
start_eigenvalue=-2e-18)
model.train(driver,
2000,
grid,
perturb,
int(8e3),
max_required_loss=1e-2,
rtol=0.01,
fraction=6,
reg_param=1e-3,
pde_param=1)
model.plot_history()
n_plot = 100
grid = torch.cat(
[grid_x.reshape(-1, 1),
grid_y.reshape(-1, 1),
grid_z.reshape(-1, 1)],
dim=1)
x = perturb(grid)
x.requires_grad = True
r = transformation[0](x)
psi = 1 / sqrt(pi) * torch.exp(-r)
model = EigenvalueProblemModel([3, 15, 15, 15, 15, 1],
torch.nn.Tanh,
compose_psi,
PDE_loss,
lr=8e-3,
betas=[0.99, 0.999],
start_eigenvalue=[-13.6 / 4, 1., 0.],
transformation=transformation,
get_energy=lambda En, En_parts: En_parts[:, 0])
model.train(driver,
2000,
grid,
perturb,
int(60e3),
max_required_loss=1e-2,
rtol=0.01,
fraction=6,
reg_param=1e-3,
pde_param=1)
model.plot_history()
def compose_psi(x, N):
f_b = 0
dt = x - x_min
psi = f_b + (1 - torch.exp(-dt)) * (1 - torch.exp(dt - x_max)) * N
return psi
def driver(index):
return -4 + (index // 2500)
grid = torch.linspace(x_min, x_max, 100).reshape(-1, 1)
model = EigenvalueProblemModel([1, 10, 10, 1], Sin, compose_psi, PDE_loss)
model.train(driver=driver,
drive_step=2500,
grid=grid,
perturb=partial(perturb1D, x_min=x_min, x_max=x_max),
epochs=int(400e3),
minibatches=1,
max_required_loss=1e-3,
fraction=3,
rtol=0.001)
model.plot_history()
large_grid = torch.linspace(x_min, x_max, 400).reshape(-1, 1)
for marker in model.eigenfunctions.keys():
y = model.get_eigenfunction(marker)(large_grid)
# x[-n_train:,0] = torch.ones_like(x[:n_train,1])
# x[0::n_train,1] = torch.zeros_like(x[::n_train,1])
# x[n_train-1::n_train,1] = torch.ones_like(x[::n_train,1])
return x
def driver(index):
return 9 + 0.25 * (index // 3000)
grid1D = torch.linspace(x_min, x_max, 10)
grid2D_x, grid2D_y = torch.meshgrid(grid1D, grid1D)
grid = torch.cat([grid2D_x.reshape(-1, 1), grid2D_y.reshape(-1, 1)], dim=1)
model = EigenvalueProblemModel([2, 20, 20, 1],
Sin,
compose_psi,
PDE_loss,
lr=1e-4,
start_eigenvalue=9.0)
model.train(driver,
3000,
grid,
perturb,
int(125e3),
max_required_loss=1e-2,
rtol=0.01,
fraction=6)
model.plot_history()
def driver(index):
return 1.3 * E_0 - 0.01 * E_0 * (index // 1000)
grid1D = torch.linspace(0, bohr_radius, 10)
grid_x, grid_y, grid_z = torch.meshgrid(grid1D, grid1D, grid1D)
grid = torch.cat(
[grid_x.reshape(-1, 1),
grid_y.reshape(-1, 1),
grid_z.reshape(-1, 1)],
dim=1)
model = EigenvalueProblemModel([3, 20, 20, 20, 1],
torch.nn.Tanh,
compose_psi,
PDE_loss,
lr=8e-3,
start_eigenvalue=1.2 * E_0,
lower_bound=0,
upper_bound=bohr_radius)
try:
model.train(driver,
1000,
grid,
perturb,
int(60e3),
max_required_loss=1e-2,
rtol=0.01,
fraction=6,
reg_param=1e-3,
pde_param=1)
except KeyboardInterrupt:
return x
def driver(index):
return 1. + .25 * (index // 1000)
grid1D = torch.linspace(x_min, x_max, 20)
grid2D_x, grid2D_y = torch.meshgrid(grid1D, grid1D)
grid = torch.cat([grid2D_x.reshape(-1, 1), grid2D_y.reshape(-1, 1)], dim=1)
model = EigenvalueProblemModel([2, 20, 20, 20, 1],
Sin,
compose_psi,
PDE_loss,
lr=8e-3,
betas=[0.9, 0.999],
start_eigenvalue=[1.2, 0.8])
model.train(driver,
2000,
grid,
perturb,
int(60e3),
max_required_loss=1e-2,
rtol=0.01,
fraction=6,
reg_param=1e-3,
pde_param=1)
model.plot_history()