def propose(self, trace, scaffold):
pnodes = scaffold.getPrincipalNodes()
currentValues = getCurrentValues(trace, pnodes)
# So the initial detach will get the gradient right
registerDeterministicLKernels(trace, scaffold, pnodes, currentValues)
rhoWeight = self.prepare(trace, scaffold,
True) # Gradient is in self.rhoDB
momenta = self.sampleMomenta(currentValues, trace.np_rng)
start_K = self.kinetic(momenta)
grad = GradientOfRegen(trace, scaffold, pnodes)
def grad_potential(values):
# The potential function we want is - log density
return [-dx for dx in grad(values)]
# Might as well save a gradient computation, since the initial
# detach does it
start_grad_pot = [-self.rhoDB.getPartial(pnode) for pnode in pnodes]
# Smashes the trace but leaves it a torus
(proposed_values, end_K) = self.evolve(grad_potential, currentValues,
start_grad_pot, momenta)
xiWeight = grad.regen(proposed_values) # Mutates the trace
# The weight arithmetic is given by the Hamiltonian being
# -weight + kinetic(momenta)
return (trace, xiWeight - rhoWeight + start_K - end_K)
def __call__(self, values):
"""Returns the gradient of the weight of regenerating along
an (implicit) scaffold starting with the given values. Smashes
the trace, but leaves it a torus. Assumes there are no delta
kernels around."""
# TODO Assert that no delta kernels are requested?
self.fixed_regen(values)
new_scaffold = constructScaffold(self.trace, [OrderedSet(self.pnodes)])
registerDeterministicLKernels(self.trace, new_scaffold, self.pnodes,
values)
(_, rhoDB) = detachAndExtract(self.trace, new_scaffold, True)
self.scaffold = new_scaffold
return [rhoDB.getPartial(pnode) for pnode in self.pnodes]
def log_joint_at(trace, args):
(scaffolders, transitions, _) = dispatch_arguments(trace, ("bogon",) + args)
if transitions > 0:
# Don't want to think about the interaction between 'each' and the
# value this is supposed to return.
assert len(scaffolders) == 1, "log_joint_at doesn't support 'each'"
scaffold = scaffolders[0].sampleIndex(trace)
pnodes = scaffold.getPrincipalNodes()
currentValues = getCurrentValues(trace, pnodes)
registerDeterministicLKernels(trace, scaffold, pnodes, currentValues,
unconditional=True)
(_rhoWeight, rhoDB) = detachAndExtract(trace, scaffold)
xiWeight = regenAndAttach(trace, scaffold, True, rhoDB, OrderedDict())
# Old state restored, don't need to do anything else
return xiWeight
else:
return 0.0
def propose(self, trace, scaffold):
pnodes = scaffold.getPrincipalNodes()
currentValues = getCurrentValues(trace, pnodes)
# So the initial detach will get the gradient right
registerDeterministicLKernels(trace, scaffold, pnodes, currentValues)
_rhoWeight = self.prepare(trace, scaffold,
True) # Gradient is in self.rhoDB
grad = GradientOfRegen(trace, scaffold, pnodes)
# Might as well save a gradient computation, since the initial
# detach does it
start_grad = [self.rhoDB.getPartial(pnode) for pnode in pnodes]
# Smashes the trace but leaves it a torus
proposed_values = self.evolve(grad, currentValues, start_grad)
_xiWeight = grad.regen(proposed_values) # Mutates the trace
return (trace, 1000) # It's MAP -- try to force acceptance
def compute_particles(self, trace, scaffold):
assertTrace(trace, scaffold)
pnodes = scaffold.getPrincipalNodes()
currentValues = getCurrentValues(trace, pnodes)
registerDeterministicLKernels(trace, scaffold, pnodes, currentValues)
rhoWeight, self.rhoDB = detachAndExtract(trace, scaffold)
xiWeights = []
xiParticles = []
allSetsOfValues = getCartesianProductOfEnumeratedValues(trace, pnodes)
for newValues in allSetsOfValues:
if newValues == currentValues:
# If there are random choices downstream, keep their current values.
# This follows the auxiliary variable method in Neal 2000,
# "Markov Chain Sampling Methods for Dirichlet Process Models"
# (Algorithm 8 with m = 1).
# Otherwise, we may target the wrong stationary distribution.
# See testEnumerativeGibbsBrushRandomness in
# test/inference_language/test_enumerative_gibbs.py for an
# example.
shouldRestore = True
omegaDB = self.rhoDB
else:
shouldRestore = False
omegaDB = OmegaDB()
xiParticle = self.copy_trace(trace)
assertTorus(scaffold)
registerDeterministicLKernels(trace, scaffold, pnodes, newValues)
xiParticles.append(xiParticle)
xiWeights.append(
regenAndAttach(xiParticle, scaffold, shouldRestore, omegaDB,
OrderedDict()))
# if shouldRestore:
# assert_almost_equal(xiWeights[-1], rhoWeight)
return (xiParticles, xiWeights)
def regen(self, values):
registerDeterministicLKernels(self.trace, self.scaffold, self.pnodes,
values)
return regenAndAttach(self.trace, self.scaffold, False, OmegaDB(),
OrderedDict())