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