Commit 0427c08e authored by Philipp Arras's avatar Philipp Arras
Browse files

Tentative fix but still not working

parent fdcda8cc
...@@ -202,8 +202,8 @@ class StochasticEnergyAdapter(Energy): ...@@ -202,8 +202,8 @@ class StochasticEnergyAdapter(Energy):
raise ValueError raise ValueError
samdom[k] = op.domain[k] samdom[k] = op.domain[k]
samdom = MultiDomain.make(samdom) samdom = MultiDomain.make(samdom)
sseq = random.spawn_sseq(n_samples) seed = int(random.current_rng().integers(0, 1000000))
context = op, position, sseq, comm, n_samples, mirror_samples, samdom context = op, position, seed, comm, n_samples, mirror_samples, samdom
noise = cls._draw_noise(*context) noise = cls._draw_noise(*context)
local_ops = [op.simplify_for_constant_input(nn)[1] for nn in noise] local_ops = [op.simplify_for_constant_input(nn)[1] for nn in noise]
n_samples = 2*n_samples if mirror_samples else n_samples n_samples = 2*n_samples if mirror_samples else n_samples
...@@ -211,8 +211,10 @@ class StochasticEnergyAdapter(Energy): ...@@ -211,8 +211,10 @@ class StochasticEnergyAdapter(Energy):
n_samples, comm, nanisinf, context, _callingfrommake=True) n_samples, comm, nanisinf, context, _callingfrommake=True)
@staticmethod @staticmethod
def _draw_noise(op, position, sseq, comm, n_samples, mirror_samples, sample_domain): def _draw_noise(op, position, seed, comm, n_samples, mirror_samples, sample_domain):
from .kl_energies import _get_lo_hi from .kl_energies import _get_lo_hi
with random.Context(seed):
sseq = random.spawn_sseq(n_samples)
noise = [] noise = []
for i in range(*_get_lo_hi(comm, n_samples)): for i in range(*_get_lo_hi(comm, n_samples)):
with random.Context(sseq[i]): with random.Context(sseq[i]):
...@@ -220,6 +222,7 @@ class StochasticEnergyAdapter(Energy): ...@@ -220,6 +222,7 @@ class StochasticEnergyAdapter(Energy):
noise.append(rnd) noise.append(rnd)
if mirror_samples: if mirror_samples:
noise.append(-rnd) noise.append(-rnd)
return noise return noise
def samples(self): def samples(self):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment