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_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])
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
simple_gaussian2 = SimpleGaussian(alpha=0.5)
jastrow2 = JastrowPade(alpha=1, beta=1)
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,
optimizer=AdamOptimizer(len(psi.parameters)),
call_backs=(psi_energies, psi_parameters),
)
from qflow.wavefunctions.nn.layers import DenseLayer
from qflow.wavefunctions.nn.activations import tanh, exponential
from qflow.hamiltonians import CoulombHarmonicOscillator
from qflow.samplers import ImportanceSampler
from qflow.statistics import compute_statistics_for_series, statistics_to_tex
from qflow.mpi import mpiprint
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)
psi.parameters = np.loadtxt("QD-parameters-dnn-regular.txt")
# Sorted
simple_gaussian2 = SimpleGaussian(alpha=0.5)
L = (P / rho)**(1 / 3)
system = np.empty((P, D))
mpiprint(f"P = {P}, D = {D}, rho = {rho}, L = {L}")
H = LennardJones(L)
layers = [
DenseLayer(P * D, 144, activation=tanh, scale_factor=0.001),
DenseLayer(144, 36, activation=tanh),
DenseLayer(36, 1, activation=exponential),
]
dnn = Dnn()
for l in layers:
dnn.add_layer(l)
mcmillian = JastrowMcMillian(5, 2.85, L)
psi_total = WavefunctionProduct(mcmillian, dnn)
psi = psi_total
sampler = HeliumSampler(system, psi, 0.5, L)
sampler.thermalize(10000)
mpiprint(f"AR: {sampler.acceptance_rate}")
# Sorted
layers2 = [
DenseLayer(P * D, 144, activation=tanh, scale_factor=0.001),
DenseLayer(144, 36, activation=tanh),
DenseLayer(36, 1, activation=exponential),
]
dnn2 = Dnn()
for l in layers2:
dnn2.add_layer(l)
mcmillian2 = JastrowMcMillian(5, 2.85, L)
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.
gamma=0, # Regularization parameter (disabled here).
optimizer=SgdOptimizer(0.1),
L = (P / rho)**(1 / 3)
system = np.empty((P, D))
mpiprint(f"P = {P}, D = {D}, rho = {rho}, L = {L}")
H = LennardJones(L)
mcmillian = JastrowMcMillian(5, 2.952, L)
layers = [
DenseLayer(2 * D, 16, activation=tanh, scale_factor=0.005),
DenseLayer(16, 1, activation=exponential),
]
dnn = Dnn()
for l in layers:
dnn.add_layer(l)
mcmillian_fixed = FixedWavefunction(mcmillian)
psi = WavefunctionProduct(mcmillian_fixed, dnn)
# psi = mcmillian
sampler = HeliumSampler(system, psi, 0.5, L)
sampler.thermalize(5000)
points = 2**int(sys.argv[2])
t0 = time.time()
psi.symmetry_metric(sampler, 500)
t1 = time.time() - t0
eta = timedelta(seconds=round(t1 / 500 * points))
mpiprint(f"Calculating symmetry - ETA {eta}")
mpiprint(f"Symmetry: {psi.symmetry_metric(sampler, points)}")
t0 = time.time()
[H.local_energy(sampler.next_configuration(), psi) / P for _ in range(500)]
t1 = time.time() - t0
H1 = CoulombHarmonicOscillator()
P, D = 8, 3 # Particles, dimensions
N = 4 # Hidden nodes
system = np.empty((P, D))
simple_gaussian = SimpleGaussian(alpha=0.5)
jastrow = JastrowOrion(beta=2, gamma=1.5)
layers = [
DenseLayer(2 * D, 32, activation=tanh, scale_factor=0.001),
DenseLayer(32, 1, activation=exponential),
]
dnn = Dnn()
for l in layers:
dnn.add_layer(l)
simple_and_jastrow = WavefunctionProduct(simple_gaussian, jastrow)
pooled_dnn = SumPooling(dnn)
psi = WavefunctionProduct(simple_and_jastrow, pooled_dnn)
psi_sampler = ImportanceSampler(system, psi, step_size=0.1)
psi_energies = EnergyCallback(samples=5000, verbose=True)
psi_symmetries = SymmetryCallback(samples=100)
psi_parameters = ParameterCallback()
train(
psi,
H1,
psi_sampler,
iters=50,
samples=1000,
gamma=0,