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,