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))