Trace NAN in NUTs

May I know if it is possible to track the NAN energy during the leapfrog steps in NUTs or HMC?

I don’t recommend this, but in the past I have inserted the following line

print("Energy:", right.energy, "  State:" right.q)

right after this code in step_methods/nuts.py

def _single_step(self, left, epsilon):
        """Perform a leapfrog step and handle error cases."""
        try:
            # `State` type
            right = self.integrator.step(epsilon, left)

This is a really quick way to get a print out of the state and energy. The q attribute refers to the canonical position coordinate for the model Hamiltonian assumed by the sampler. If you want to print out the momenta variables as well (i.e. to see if too much momentum is sending you into a bad region), you could print out right.p as well.

1 Like