def test_laplacian(X, alpha, beta):
psi = SimpleGaussian(alpha, beta)
laplacian = psi.laplacian(X) * psi(X)
assert_close(
np.trace(hessian(psi_np)(X, alpha, beta).reshape(X.size, X.size)),
laplacian)
def test_gradient(X, alpha, beta):
psi = SimpleGaussian(alpha, beta)
gradient = psi.gradient(X)
assert len(gradient) == 2
assert gradient[1] == 0 # Beta fixed.
assert_close(grad(psi_np, 1)(X, alpha, beta), gradient[0] * psi(X))
def test_set_parameters_is_noop_in_wavefunction_product():
orig1 = SimpleGaussian(0.4)
fixed = FixedWavefunction(orig1)
orig2 = SimpleGaussian(0.6)
prod = WavefunctionProduct(fixed, orig2)
# Only the params related to orig2 should change after this.
# Also, orig2's beta should not change because it is const by def.
prod.parameters = [0.5, 1.1, 123, 321]
np.testing.assert_array_equal(prod.parameters, [0.4, 1, 123, 1])
def test_consistent_with_original():
orig = SimpleGaussian()
fixed = FixedWavefunction(orig)
for _ in range(100):
n, d = np.random.randint(100), np.random.randint(1, 3 + 1)
s = np.random.randn(n, d)
assert np.isclose(orig(s), fixed(s))
assert np.isclose(orig.laplacian(s), fixed.laplacian(s))
np.testing.assert_allclose(orig.drift_force(s), fixed.drift_force(s))
np.testing.assert_array_equal(orig.parameters, fixed.parameters)
# Gradient is special, should return zeros for each parameter as they are fixed.
np.testing.assert_array_equal(orig.gradient(s) * 0, fixed.gradient(s))
def test_ideal_harmonic_oscillator():
omega, N, D = 1, 5, 3
H = HarmonicOscillator(omega_ho=omega)
psi = SimpleGaussian(alpha=0.5 * omega)
sampler = ImportanceSampler(np.empty((N, D)), psi, 0.1)
sampler.thermalize(10000)
# Energy:
E = compute_statistics_for_series(H.local_energy_array(
sampler, psi, 2**10))
assert N * D * omega / 2 == E["mean"], "Energy should be exactly equal."
assert E["var"] < 1e-16
# Squared radius:
r2 = compute_statistics_for_series(H.mean_squared_radius_array(
sampler, 2**16),
method="blocking")
print(r2)
assert (abs(3 / omega / 2 - r2["mean"]) <
1.96 * r2["sem"]), "Should be withing 95% CI of analytic result."
# Radius:
r = compute_statistics_for_series(H.mean_radius_array(sampler, 2**16),
method="blocking")
print(r)
assert (abs(2 / np.sqrt(np.pi * omega) - r["mean"]) <
1.96 * r["sem"]), "Should be withing 95% CI of analytic result."
def test_gradient(X, alpha):
psi = SimpleGaussian(alpha)
pool = SumPooling(psi)
assert_close(pool.gradient(X)[1], 0)
assert_close([grad(sum_pool_np, 1)(X, alpha) / sum_pool_np(X, alpha)],
pool.gradient(X)[:1])
def test_set_parameters_is_noop():
orig = SimpleGaussian()
fixed = FixedWavefunction(orig)
# This should do noting to the parameters.
fixed.parameters = orig.parameters + 1
np.testing.assert_array_equal(orig.parameters, fixed.parameters)
def test_works_on_simple_gaussians():
"""
If
Psi = SimpleGaussian(a) * SimpleGaussian(b)
then we have by the form of SimpleGaussian the follwing
Psi = SimpleGaussian(a + b)
for any a and b.
"""
for _ in range(1000):
alpha1, alpha2 = np.random.rand(2)
psi1 = SimpleGaussian(alpha1)
psi2 = SimpleGaussian(alpha2)
psi_expected = SimpleGaussian(alpha1 + alpha2)
psi_prod = WavefunctionProduct(psi1, psi2)
n, d = np.random.randint(100), np.random.randint(1, 3 + 1)
s = np.random.randn(n, d)
assert np.isclose(psi_expected(s), psi_prod(s))
assert np.isclose(psi_expected.laplacian(s), psi_prod.laplacian(s))
np.testing.assert_allclose(psi_expected.drift_force(s), psi_prod.drift_force(s))
# The gradient will be slightly different. Both psi1 and psi2 should give the same gradient, as it is
# independent of alpha. However, the product state will give the gradients from both psi1 and psi2,
# which means it will have two sets of each number.
# __This is the expected behaviour.__
expected_grad = psi_expected.gradient(s)
np.testing.assert_allclose(
np.concatenate([expected_grad] * 2), psi_prod.gradient(s)
)
def test_laplacian(X, alpha):
psi = SimpleGaussian(alpha)
pool = SumPooling(psi)
expected = np.trace(
hessian(sum_pool_np)(X, alpha).reshape(X.size, X.size)) / sum_pool_np(
X, alpha)
assert_close(expected, pool.laplacian(X))
import numpy as np
import matplotlib.pyplot as plt
import matplotlib2tikz
from qflow.wavefunctions import SimpleGaussian
from qflow.hamiltonians import HarmonicOscillator
from qflow.samplers import ImportanceSampler
from qflow.optimizers import SgdOptimizer, AdamOptimizer
from qflow.training import EnergyCallback, train
from qflow.mpi import master_rank
N, D = 1, 1
system = np.empty((N, D))
H = HarmonicOscillator(omega_ho=1)
psi = SimpleGaussian(0.8)
org_params = psi.parameters[:]
sampler = ImportanceSampler(system, psi, 0.5)
labels = [
r"Sgd($\eta=0.1$)",
r"Sgd($\eta=0.050$)",
r"Sgd($\eta=0.025$)",
r"Sgd($\eta=0.010$)",
r"Adam($\eta=0.05,\beta_1=0.9$)",
r"Adam($\eta=0.05,\beta_1=0.7$)",
]
optimizers = [
SgdOptimizer(0.1),
SgdOptimizer(0.05),
SgdOptimizer(0.025),
SgdOptimizer(0.01),
_, sax = plt.subplots()
sax.semilogx(symmetries, label=r"$S(\psi_{DNN})$")
sax.set_ylabel("Symmetry")
sax.set_xlabel(r"% of training")
sax.legend(loc="lower right")
matplotlib2tikz.save(__file__ + ".symmetry.tex")
P, D = 2, 2 # Particles, dimensions
system = np.empty((P, D))
H = CoulombHarmonicOscillator()
# Wave functions:
simple_gaussian = SimpleGaussian(alpha=0.5)
jastrow = JastrowPade(alpha=1, beta=1)
simple_and_jastrow = WavefunctionProduct(simple_gaussian, jastrow)
layers = [
DenseLayer(P * D, 32, activation=tanh, scale_factor=0.001),
DenseLayer(32, 16, activation=tanh),
DenseLayer(16, 1, activation=exponential),
]
dnn = Dnn()
for l in layers:
dnn.add_layer(l)
psi = WavefunctionProduct(simple_and_jastrow, dnn)
psi_sampler = ImportanceSampler(system, psi, step_size=0.1)
# Sorted
def test_gaussian_is_symmetric(P, D, alpha):
psi = SimpleGaussian(alpha)
sampler = MetropolisSampler(np.empty((P, D)), psi, 0.1)
s = psi.symmetry_metric(sampler, 10)
assert np.isclose(1, s)
eax.plot(energies, label=r"$\langle E_L\rangle$ [a.u]")
eax.set_xlabel(r"% of training")
eax.axhline(y=3, label="Exact", linestyle="--", color="k", alpha=0.5)
eax.legend()
pax.plot(np.asarray(parameters)[:, [0, 3]])
pax.set_xlabel(r"% of training")
pax.legend([r"$\alpha_G$", r"$\beta_{PJ}$"])
matplotlib2tikz.save(__file__ + ".tex")
P, D = 2, 2 # Particles, dimensions
system = np.empty((P, D))
H = CoulombHarmonicOscillator()
simple_gaussian = SimpleGaussian(alpha=0.5)
jastrow = JastrowPade(alpha=1, beta=1)
psi = WavefunctionProduct(simple_gaussian, jastrow)
psi_sampler = ImportanceSampler(system, psi, step_size=0.1)
psi_simple_sampler = ImportanceSampler(system, simple_gaussian, step_size=0.1)
psi_energies = EnergyCallback(samples=100000)
psi_parameters = ParameterCallback()
train(
psi,
H,
psi_sampler,
iters=2000,
samples=1000,
gamma=0,
from qflow.wavefunctions.nn.layers import DenseLayer
from qflow.wavefunctions.nn.activations import (
sigmoid,
tanh,
relu,
identity,
exponential,
)
from qflow import DistanceCache
small_system = np.zeros((2, 2))
large_system = np.zeros((50, 3))
samples = 10000
H0 = HarmonicOscillator()
psi0 = SimpleGaussian(0.5)
layers = [
DenseLayer(50 * 3, 32, activation=tanh, scale_factor=0.001),
DenseLayer(32, 16, activation=tanh),
DenseLayer(16, 1, activation=exponential),
]
dnn = Dnn()
for l in layers:
dnn.add_layer(l)
def local_energy_gradient(H, psi, sampler, samples):
return H.local_energy_gradient(sampler, psi, samples)
@pytest.mark.benchmark(group="evaluation", warmup=True)
import numpy as np import matplotlib.pyplot as plt import matplotlib2tikz from qflow.wavefunctions import SimpleGaussian from qflow.hamiltonians import HarmonicOscillator from qflow.samplers import ImportanceSampler from qflow.mpi import master_rank N, D = 10, 3 system = np.empty((N, D)) H = HarmonicOscillator(omega_ho=1) psi = SimpleGaussian(alpha=0.5) sampler = ImportanceSampler(system, psi, 0.1) sampler.thermalize(10000) max_r = 1.5 samples = 2**25 n_bins = 20 r = np.linspace(0, max_r, n_bins) bins = H.twobodydensity(sampler, n_bins, max_r, samples) rho = bins / np.trapz(np.trapz(bins, x=r), x=r) rho /= np.max(rho) rx, ry = np.meshgrid(r, r) exact = np.exp(-rx**2 - ry**2) exact /= np.trapz(np.trapz(exact, x=r), x=r) exact /= np.max(exact) if master_rank():
import numpy as np
from qflow.wavefunctions import SimpleGaussian
from qflow.hamiltonians import HarmonicOscillator
from qflow.samplers import ImportanceSampler
from qflow.statistics import compute_statistics_for_series, statistics_to_tex
from qflow.mpi import mpiprint
N, D = 100, 3
system = np.empty((N, D))
H = HarmonicOscillator(omega_ho=1)
psi_opt = SimpleGaussian(alpha=0.5)
psi_nopt = SimpleGaussian(alpha=0.51)
sampler_opt = ImportanceSampler(system, psi_opt, 0.1)
sampler_nopt = ImportanceSampler(system, psi_nopt, 0.1)
sampler_opt.thermalize(10000)
sampler_nopt.thermalize(10000)
samples = 2**23
stats = [
compute_statistics_for_series(
H.local_energy_array(sampler_opt, psi_opt, 100) / N),
compute_statistics_for_series(
H.local_energy_array(sampler_nopt, psi_nopt, samples) / N,
method="blocking"),
]
labels = [r"$\alpha_G = 0.5$", r"$\alpha_G=0.51$"]
from qflow.wavefunctions import SimpleGaussian, JastrowPade, WavefunctionProduct
from qflow.hamiltonians import CoulombHarmonicOscillator
from qflow.samplers import ImportanceSampler
from qflow.optimizers import SgdOptimizer
from qflow.training import train, EnergyCallback, ParameterCallback
from qflow.statistics import compute_statistics_for_series
from qflow.mpi import mpiprint, master_rank
P, D = 2, 2 # Particles, dimensions
# Define Hamiltonian:
H = CoulombHarmonicOscillator(omega_ho=1)
# Define trial wave function:
gaussian = SimpleGaussian(alpha=0.8)
jastrow = JastrowPade(alpha=1, beta=1)
psi = WavefunctionProduct(gaussian, jastrow)
# Set up sampling strategy:
sampler = ImportanceSampler(np.empty((P, D)), psi, step_size=0.1)
# Train wave function:
training_energies = EnergyCallback(samples=100000)
training_params = ParameterCallback()
train(
psi,
H,
sampler,
iters=150, # Optimization steps.
samples=1000, # MC cycles per optimization step.
def test_drift_force(X, alpha):
psi = SimpleGaussian(alpha)
pool = SumPooling(psi)
expected = 2 * grad(sum_pool_np, 0)(X, alpha) / sum_pool_np(X, alpha)
assert_close(expected.ravel(), pool.drift_force(X))
def test_eval(X, alpha):
psi = SimpleGaussian(alpha)
pool = SumPooling(psi)
numpy.testing.assert_allclose(sum_pool_np(X, alpha), pool(X))
import numpy as np
import matplotlib.pyplot as plt
import matplotlib2tikz
from qflow.wavefunctions import SimpleGaussian
from qflow.hamiltonians import HarmonicOscillator
from qflow.samplers import ImportanceSampler
from qflow.optimizers import SgdOptimizer
from qflow.training import train, EnergyCallback, ParameterCallback
N, D = 10, 3
system = np.empty((N, D))
H = HarmonicOscillator(omega_ho=1)
psi_G = SimpleGaussian(alpha=0.3)
isampler_G = ImportanceSampler(system, psi_G, 0.1)
isampler_G.thermalize(10000)
E_training = EnergyCallback(samples=50000, verbose=True)
G_training = ParameterCallback()
train(
psi_G,
H,
isampler_G,
iters=100,
samples=1000,
gamma=0,
optimizer=SgdOptimizer(0.001),
call_backs=[E_training, G_training],
)
E_training = np.asarray(E_training) / N
def test_eval(X, alpha, beta):
psi = SimpleGaussian(alpha, beta)
assert_close(psi_np(X, alpha, beta), psi(X))
def test_drift_force(X, alpha, beta):
psi = SimpleGaussian(alpha, beta)
drift = 0.5 * psi.drift_force(X) * psi(X)
assert len(drift) == X.size
assert_close(grad(psi_np, 0)(X, alpha, beta).ravel(), drift)