def binary_mse_multinomial(repeats, niterations, f_range, beta, nstates=2):
"""
Function to compute the mean-squared error of the SAMS binary update scheme when drawing samples from
the multinomial distribution. Over many repeats, target free energies will be drawn randomly and uniformly
from a specified range and SAMS will adapt to those free energies.
Parameters
----------
repeats: int
The number of repeats with which to draw target free energies and run SAMS
ninterations:
The number of state samples generated and SAMS adaptive steps
f_range: float
The interval over which target free energies will be drawn
beta: float
The exponent for the SAMS burn-in stage
nstates: int
The number of states and target free energies.
"""
binary_aggregate_msd = np.zeros((repeats, niterations))
for r in range(repeats):
f_true = np.random.uniform(low=-f_range / 2.0, high=f_range / 2.0, size=nstates)
f_true -= f_true[0]
generator = IndependentMultinomialSamper(free_energies=f_true)
adaptor = SAMSAdaptor(nstates=nstates, beta=beta)
for i in range(niterations):
noisy = generator.step()
state = np.where(noisy != 0)[0][0]
z = -adaptor.update(state=state, noisy_observation=noisy, histogram=generator.histogram)
generator.zetas = z
binary_aggregate_msd[r, i] = np.mean((f_true - z) ** 2)
return binary_aggregate_msd
def test_slow_gain(self):
"""
Ensure the correct gain factor is correctly calculated in the absence of two stage procedure.
"""
# The system and SAMS parameters
nstates = 10 # The number of discrete states in the system
target_weights = np.repeat(1.0 / nstates,
nstates) # The target probabilities
initial_zeta = np.repeat(
1.0, nstates) # Initial estimate for the free energies
initial_zeta -= initial_zeta[0]
adaptor = SAMSAdaptor(nstates=nstates,
two_stage=False,
zetas=initial_zeta,
target_weights=target_weights)
# Generate a fake sample for the noisy variable:
noisy = np.zeros(nstates)
state = 3 # The imagined state of the sampler
noisy[state] = 1.0
new_zetas = adaptor.update(state=state, noisy_observation=noisy)
true_new_zeta = initial_zeta[state] + (
noisy[state] / target_weights[state]) / adaptor.time
assert new_zetas[state] == true_new_zeta
def test_sams_clock(self):
"""
Assess that the adaptor accurately logs the number of times it was called and its iteration stage. The SAMS
internal clock is used to calculate the gain. Samples will be generated with numpy's multinomial sampler.
"""
# Initialize the SAMS adaptor
nstates = 10
adaptor = SAMSAdaptor(nstates=nstates)
# The probabilities for the discrete states
probs = np.repeat(1.0 / nstates, nstates)
# Iterating the sampler and adaptor for few steps
niterations = 10
for i in range(niterations):
noisy = np.random.multinomial(1, probs)
state = np.where(noisy != 0)[0][0]
adaptor.update(state=state, noisy_observation=noisy)
assert adaptor.time == niterations
def test_burnin_termination(self):
"""
Tests whether the stopping criterion for the burn-in period is correctly implemented.
"""
nstates = 10
target_weights = np.repeat(1.0 / nstates, nstates)
adaptor = SAMSAdaptor(nstates=nstates,
two_stage=True,
flat_hist=0.2,
target_weights=target_weights)
# Creating a histogram that matches the target weights. The burn-in should stop.
histogram = 100 * target_weights
noisy = np.zeros(nstates)
state = 3 # The imagined state of the sampler
noisy[state] = 1.0
adaptor.update(state=state,
noisy_observation=noisy,
histogram=histogram)
assert adaptor.burnin == False
def test_burnin_continuation(self):
"""
Tests whether the burn-in does not stop prematurely
"""
nstates = 10
target_weights = np.repeat(1.0 / nstates, nstates)
adaptor = SAMSAdaptor(nstates=nstates,
two_stage=True,
flat_hist=0.2,
target_weights=target_weights)
# Creating a histogram that is far from 'flat' and therefore should not stop the burn-in phase
histogram = np.arange(nstates)
# Creating a fake sample
noisy = np.zeros(nstates)
state = 3 # The imagined state of the sampler
noisy[state] = 1.0
adaptor.update(state=state,
noisy_observation=noisy,
histogram=histogram)
assert adaptor.burnin == True
def rb_mse_gaussian(sigmas, niterations, nmoves=1, save_freq=1, beta=0.6, flat_hist=0.2):
"""
Function to compute the mean-squared error from the Rao-Blackwellized update scheme in when sampling
states with GaussianMixtureSampler.
Parameters
----------
sigmas: numpy array
The standard deviations of the Gaussians that are centered on zero
niterations: int
The number of iterations of mixture sampling and SAMS updates that will be performed
nmoves: int
The number of moves from the Gaussian Gibbs sampler. One position step and state step
constitute one move.
save_freq: int
The frequency with which to save the state in the Gaussian mixture sampler, used to decorrelate trajectory
beta: float
The exponent in the burn-in phase of the two stage SAMS update scheme.
flat_hist: float
The average fractional difference between the target weights of the mixture and count frequency
Returns
-------
mse: numpy array
the mean-squared error of the SAMS estimate for each iteration
"""
generator = GaussianMixtureSampler(sigmas=sigmas)
nstates = len(sigmas)
adaptor = SAMSAdaptor(nstates=nstates, beta=beta, flat_hist=flat_hist)
# The target free energy
f_true = -np.log(sigmas)
f_true = f_true - f_true[0]
mse = np.zeros((niterations))
for i in range(niterations):
generator.step(nmoves, save_freq)
state = generator.state
noisy = generator.weights
z = -adaptor.update(state=state, noisy_observation=noisy, histogram=generator.histogram)
generator.zetas = z
mse[i] = np.mean((f_true[1:] - z[1:]) ** 2)
return mse
def test_fast_gain(self):
"""
Ensures the gain during the burn-in stage is properly calculated.
"""
# The system and SAMS parameters
nstates = 10 # The number of discrete states in the system
beta = 0.7 # Exponent of gain factor during burn-in phase.
target_weights = np.repeat(1.0 / nstates,
nstates) # The target probabilities
initial_zeta = np.repeat(
1.0, nstates) # Initial estimate for the free energies
initial_zeta -= initial_zeta[0]
adaptor = SAMSAdaptor(nstates=nstates,
two_stage=True,
zetas=initial_zeta,
target_weights=target_weights,
beta=beta)
# Generate a fake sample for the noisy variable:
noisy = np.zeros(nstates)
state = 3 # The imagined state of the sampler
noisy[state] = 1.0
new_zetas = adaptor.update(state=state,
noisy_observation=noisy,
histogram=np.arange(nstates))
if target_weights[state] < adaptor.time**(-beta):
true_new_zeta = initial_zeta[state] + target_weights[state] * (
noisy[state] / target_weights[state])
else:
true_new_zeta = initial_zeta[state] + adaptor.time**(-beta) * (
noisy[state] / target_weights[state])
assert new_zetas[state] == true_new_zeta
nwater, ncation, nanion = mcdriver.get_identity_counts()
penalty = gen_penalty(ncation, bias, args.saltmax)
# Propagate configurations and salt concentrations
langevin.step(args.steps)
mcdriver.attempt_identity_swap(context,
penalty=penalty,
saltmax=args.saltmax)
# Update SAMS estimator
nwater, ncation, nanion = mcdriver.get_identity_counts()
noisy = np.zeros(args.saltmax + 1)
noisy[ncation] = 1
state_counts[ncation] += 1
bias = adaptor.update(state=ncation,
noisy_observation=noisy,
histogram=state_counts)
# Save data
if iteration % args.save_freq == 0:
# Record the simulation data
Record.record_netcdf(ncfile,
context,
mcdriver,
k,
attempt=0,
sync=False)
ncfile.groups['Sample state data']['sams bias'][k, :] = bias
ncfile.sync()
if args.save_configs:
# Record the simulation configurations