def setUp(self): auto_np.random.seed(1234) self.nn = Dnn() self.input_layer = DenseLayer(2, 3, activation=sigmoid) self.middle_layer = DenseLayer(3, 4, activation=relu) self.output_layer = DenseLayer(4, 1) self.nn.add_layer(self.input_layer) self.nn.add_layer(self.middle_layer) self.nn.add_layer(self.output_layer) self.W1 = self.nn.layers[0].weights self.b1 = self.nn.layers[0].biases self.W2 = self.nn.layers[1].weights self.b2 = self.nn.layers[1].biases self.W3 = self.nn.layers[2].weights self.b3 = self.nn.layers[2].biases self.params = [self.W1, self.b1, self.W2, self.b2, self.W3, self.b3] def f(x, w1, b1, w2, b2, w3, b3): z1 = x @ w1 + b1 z2 = sigmoid_np(z1) @ w2 + b2 z3 = relu_np(z2) @ w3 + b3 return z3 self.f_np = f
def test_symmetry(P, D): dnn = Dnn() layer = DenseLayer(P * D, 1) dnn.add_layer(layer) sdnn = InputSorter(dnn) # Random system, sorted by distance to origin. s = np.random.randn(P, D) s = s[np.argsort([row.dot(row) for row in s])] # Limit the number of permutations, for computational feasibility. for _, p in zip(range(100), permutations(s)): p = np.asarray(p) # Permutation invariant: assert sdnn(s) == sdnn(p) # Values should be equal to dnn applied to the pre-sorted system. assert sdnn(p) == dnn(s) assert sdnn.laplacian(p) == dnn.laplacian(s) assert all(sdnn.gradient(p) == dnn.gradient(s)) # Drift force should have equal values, but permuted according to the permutation p. # This because the drift force depends on which positions are which. Values should # still be the same, so we only allow for sequence particle-wise permutations. sdnn_drift, dnn_drift = ( sdnn.drift_force(p).reshape(s.shape), dnn.drift_force(s).reshape(s.shape), ) sdnn_drift.sort(axis=0) dnn_drift.sort(axis=0) assert (sdnn_drift == dnn_drift).all()
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) simple_and_jastrow2 = WavefunctionProduct(simple_gaussian2, jastrow2) layers2 = [
matplotlib2tikz.save(__file__ + ".tex") os.makedirs("logfiles", exist_ok=True) rho = 0.365 / (2.556)**3 # Å^-3 P, D = 32, 3 # Particles, dimensions 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 = [
sigmoid, tanh, ) from qflow.wavefunctions.nn.layers import DenseLayer rho = 0.365 / (2.556)**3 # Å^-3 P, D = 4, 3 # Particles, dimensions 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