def test_autograd_trace_hess_batchdims():
x1 = torch.ones(1, 3)
lap = autograd_trace_hessian(func_quadratic, x1)
assert lap.shape == torch.Size([1])
x1 = torch.ones(5, 3)
lap = autograd_trace_hessian(func_quadratic, x1)
assert lap.shape == torch.Size([5])
x1 = torch.ones(10, 5, 3)
lap = autograd_trace_hessian(func_quadratic, x1)
assert lap.shape == torch.Size([10, 5])
def test_autograd_trace_hess_batchdims():
x1 = torch.ones(1, 3)
lap, grad = autograd_trace_hessian(func_quadratic, x1, return_grad=True)
assert lap.shape == torch.Size([1])
assert grad.shape == torch.Size([1, 3])
x1 = torch.ones(5, 3)
lap, grad = autograd_trace_hessian(func_quadratic, x1, return_grad=True)
assert lap.shape == torch.Size([5])
assert grad.shape == torch.Size([5, 3])
x1 = torch.ones(10, 5, 3)
lap, grad = autograd_trace_hessian(func_quadratic, x1, return_grad=True)
assert lap.shape == torch.Size([10, 5])
assert grad.shape == torch.Size([10, 5, 3])
def test_hydrogen_ansatz_autograd(x):
x = torch.tensor(x)
#
func = HydrogenTrialWavefunction(torch.tensor(1.0))
true_lap = (func.alpha**2 * torch.exp(-func.alpha * x) *
(-2.0 + func.alpha * x)).squeeze(dim=-1)
#
auto_lap = autograd_trace_hessian(func.hydro_ansatz_sup, x)
assert true_lap.shape == auto_lap.shape
assert torch.isclose(true_lap, auto_lap).all()
def auto_hamiltonian_generator_atoms(ansatz, N_bodies, config): #Unitless for now.
kinetic_source = autograd_trace_hessian( ansatz.wave,config, return_grad =True)
kinetic_energy_0 = -0.5*kinetic_source[0]
kinetic_energy_1 = 0
for i in range(N_bodies):
r = (kinetic_source[1][..., i])/config[...,i]
kinetic_energy_1 = kinetic_energy_1-r
kinetic_total = kinetic_energy_0 + kinetic_energy_1
potential_nucleous_electron = 0
for i in range(N_bodies):
potential_nucleous_electron = potential_nucleous_electron-(N_bodies/(config[...,i].squeeze(dim=-1)))*ansatz.wave(config)
potential_electron_electron = 0
k = 1
for i in range(N_bodies):
for j in range(i+1,N_bodies):
potential_electron_electron = potential_electron_electron + (ansatz.wave(config))/(torch.sqrt(config[...,i]**2+config[...,j]**2+torch.abs(config[...,i])*torch.abs(config[...,j])*torch.cos(config[...,N_bodies -1+k])))
k+=1
energy_total = kinetic_total+potential_nucleous_electron+potential_electron_electron
return energy_total
def test_autograd_trace_hess_exp(x):
xtens = torch.tensor(x)
out = func_exp(xtens) # deriv of exp is exp
assert torch.isclose(autograd_trace_hessian(func_exp, xtens), out).all()
def test_autograd_trace_hess_quadratic(x):
assert autograd_trace_hessian(func_quadratic, torch.tensor(x)) == 60.0
def test_offdiag_laplacian(x):
x = torch.tensor(x)
assert np.isclose(autograd_trace_hessian(func_offdiag, x)[0].item(), 0.0)
def test_full_hess_equivalent(x):
x = torch.tensor(x)
x2 = x.clone()
assert np.isclose(
autograd_trace_hessian(func_offdiag, x)[0].item(),
autograd_trace_hessian_full(func_offdiag, x2).item())
def test_autograd_trace_cubic():
assert autograd_trace_hessian(func_cubic, torch.ones(1, 3)) == 180.0
def test_runs_autograd_trace_cubic():
autograd_trace_hessian(func_cubic, torch.ones(1, 3))
def test_autograd_trace_hess_sin(x):
input = torch.tensor([[x]])
lap = autograd_trace_hessian(func_sin, input)
assert torch.isclose(lap, -torch.sin(input))
def local_energy(self,x):
# return -0.5*autograd_trace_hessian(self.slater_ansatz_2_particle_in_box,x)/self.slater_ansatz_2_particle_in_box(x)
return -autograd_trace_hessian(self.slater_ansatz_2_particle_in_box,x,return_grad = False)/self.slater_ansatz_2_particle_in_box(x)
def local_energy(self, x):
return (-(1.0/x.squeeze(dim=-1))-(0.5*autograd_trace_hessian(self.hydro_ansatz_sup,x)/(self.hydro_ansatz_sup(x))))
def local_energy(self, x):
return ((x**2).squeeze(dim=-1)-(autograd_trace_hessian(self.harmoni_ansatz_sup,x)/(self.harmoni_ansatz_sup(x))))
import torch
import numpy as np
from hypothesis import given
from hypothesis.extra.numpy import arrays, array_shapes
from hypothesis.strategies import floats
from qmc.tracehess import autograd_trace_hessian, gradient_f, hessian_f
# Gradient Function
#def func_cubic(x):
# return 10.0 * torch.sum(x * x * x * x, dim=-1)
#print(autograd_trace_hessian(func_cubic, y))
self_alpha=1
y = torch.ones(1,1)
def func_exp(x):
return torch.exp(x).prod(dim=-1)
print(autograd_trace_hessian(func_exp, y))
def hydro_ansatz_sup(x):
return self_alpha*x*torch.exp(-self_alpha*x)
print(autograd_trace_hessian(hydro_ansatz_sup, y)/hydro_ansatz_sup(y))