Make sure we are handling NPT correctly with MonteCarloBarostat
It appears that we aren't correctly treating solvated systems by adding a `MonteCarloBarostat` to them. While we _don't_ want to include volume change in NCMC (though this might be OK when GHMC is used), we definitely want to add a barostat for solvated systems, since otherwise the default water densities can be awful, making NCMC much more difficult than it should be.

diff --git a/devtools/conda-recipe/meta.yaml b/devtools/conda-recipe/meta.yaml
index 22208ab1..de0c9dc4 100644
--- a/devtools/conda-recipe/meta.yaml
+++ b/devtools/conda-recipe/meta.yaml
@@ -18,10 +18,10 @@ requirements:
     - numexpr
     - autograd
     - pymbar
-    - openmm
+    - openmm >=7.0.1
     - parmed
     - openmoltools
-    - alchemy
+    - alchemy >=1.2.0
     - openmmtools
     - numba
     - netcdf4
@@ -37,10 +37,10 @@ requirements:
     - numexpr
     - autograd
     - pymbar
-    - openmm
+    - openmm >=7.0.1
     - parmed
     - openmoltools
-    - alchemy
+    - alchemy >=1.2.0
     - openmmtools
     - numba
     - netcdf4
diff --git a/perses/analysis/analysis.py b/perses/analysis/analysis.py
index e64643a1..4ec8449c 100644
--- a/perses/analysis/analysis.py
+++ b/perses/analysis/analysis.py
@@ -23,7 +23,7 @@ import json
 
 import matplotlib as mpl
 mpl.use('Agg')
-import seaborn
+import seaborn as sns
 
 from matplotlib.backends.backend_pdf import PdfPages
 import matplotlib.pyplot as plt
@@ -87,7 +87,9 @@ class Analysis(object):
                     except Exception as e:
                         pass
 
-                if len(work) > 0:
+                def plot_work_trajectories(pdf, work, title=""):
+                    """Generate figures for the specified switching legs.
+                    """
                     plt.figure(figsize=(12, 8))
 
                     nrows = 2
@@ -113,7 +115,7 @@ class Analysis(object):
                         # Label plot
                         if row == 1: plt.xlabel('steps')
                         plt.ylabel('work / kT')
-                        plt.title("NCMC in environment '%s' : %s" % (envname, direction))
+                        plt.title("%s NCMC in environment '%s' : %s" % (title, envname, direction))
                         plt.legend(['average work', 'NCMC attempts'])
 
                         #
@@ -124,9 +126,10 @@ class Analysis(object):
                         plt.subplot2grid((nrows,ncols), (row, col), colspan=workcols)
 
                         # Plot average work distribution in think solid line
-                        nbins = 40
+                        #nbins = 40
                         workvals = work[direction][:-1,-1]
-                        plt.hist(workvals, nbins)
+                        #plt.hist(workvals, nbins)
+                        sns.distplot(workvals, rug=True)
                         # Adjust axes to eliminate large-magnitude outliers (keep 98% of data in-range)
                         #worklim = np.percentile(workvals, 98)
                         #oldaxis = plt.axis()
@@ -135,6 +138,15 @@ class Analysis(object):
                         if row == 1: plt.xlabel('work / kT')
                         plt.title("total %s work" % direction)
 
-
                     pdf.savefig()  # saves the current figure into a pdf page
                     plt.close()
+
+                if len(work) > 0:
+                    # Plot work for all chemical transformations.
+                    plot_work_trajectories(pdf, work, title='(all transformations)')
+
+                    # Plot work separated out for each chemical transformation
+                    #[niterations, nsteps] = work.shape
+                    #transformations = dict()
+                    #for iteration in range(niterations):
+                    #    plot_work_trajectories(pdf, work, title='(all transformations)')
diff --git a/perses/annihilation/ncmc_switching.py b/perses/annihilation/ncmc_switching.py
index f4799ae4..f7118f2f 100644
--- a/perses/annihilation/ncmc_switching.py
+++ b/perses/annihilation/ncmc_switching.py
@@ -6,10 +6,11 @@ import traceback
 from simtk import openmm, unit
 from openmmtools.integrators import GHMCIntegrator
 from perses.storage import NetCDFStorageView
+from perses.tests.utils import quantity_is_finite
 
 default_functions = {
-    'lambda_sterics' : 'lambda',
-    'lambda_electrostatics' : 'lambda',
+    'lambda_sterics' : '2*lambda * step(0.5 - lambda) + (1.0 - step(0.5 - lambda))',
+    'lambda_electrostatics' : '2*(lambda - 0.5) * step(lambda - 0.5)',
     'lambda_bonds' : '0.9*lambda + 0.1', # don't fully soften bonds
     'lambda_angles' : '0.9*lambda + 0.1', # don't fully soften angles
     'lambda_torsions' : 'lambda'
@@ -45,7 +46,7 @@ class NCMCEngine(object):
 
     """
 
-    def __init__(self, temperature=default_temperature, functions=None, nsteps=default_nsteps, steps_per_propagation=default_steps_per_propagation, timestep=default_timestep, constraint_tolerance=None, platform=None, write_pdb_interval=None, integrator_type='GHMC', storage=None):
+    def __init__(self, temperature=default_temperature, functions=None, nsteps=default_nsteps, steps_per_propagation=default_steps_per_propagation, timestep=default_timestep, constraint_tolerance=None, platform=None, write_ncmc_interval=None, integrator_type='GHMC', storage=None, verbose=False):
         """
         This is the base class for NCMC switching between two different systems.
 
@@ -66,13 +67,15 @@ class NCMCEngine(object):
             If not None, this relative constraint tolerance is used for position and velocity constraints.
         platform : simtk.openmm.Platform, optional, default=None
             If specified, the platform to use for OpenMM simulations.
-        write_pdb_interval : int, optional, default=None
-            If a positive integer is specified, a PDB frame will be written with the specified interval on NCMC switching,
-            with a different PDB file generated for each attempt.
+        write_ncmc_interval : int, optional, default=None
+            If a positive integer is specified, a snapshot frame will be written to storage with the specified interval on NCMC switching.
+            'storage' must also be specified.
         integrator_type : str, optional, default='GHMC'
             NCMC internal integrator type ['GHMC', 'VV']
         storage : NetCDFStorageView, optional, default=None
             If specified, write data using this class.
+        verbose : bool, optional, default=False
+            If True, print debug information.
         """
         # Handle some defaults.
         if functions == None:
@@ -92,14 +95,17 @@ class NCMCEngine(object):
         self.platform = platform
         self.integrator_type = integrator_type
         self.steps_per_propagation = steps_per_propagation
+        self.verbose = verbose
 
-        self.write_pdb_interval = write_pdb_interval
+        if steps_per_propagation != 1:
+            raise Exception('steps_per_propagation must be 1 until CustomIntegrator is debugged')
 
         self.nattempted = 0
 
         self._storage = None
         if storage is not None:
             self._storage = NetCDFStorageView(storage, modname=self.__class__.__name__)
+        self.write_ncmc_interval = write_ncmc_interval
 
     @property
     def beta(self):
@@ -247,13 +253,9 @@ class NCMCEngine(object):
         else:
             raise Exception("direction must be one of ['delete', 'insert']; found '%s' instead" % direction)
 
-        # DEBUG
-        #print('alchemical atoms:')
-        #print(alchemical_atoms)
-
         # Create an alchemical factory.
         from alchemy import AbsoluteAlchemicalFactory
-        alchemical_factory = AbsoluteAlchemicalFactory(unmodified_system, ligand_atoms=alchemical_atoms, annihilate_electrostatics=True, annihilate_sterics=True, alchemical_bonds=None, alchemical_angles=None)
+        alchemical_factory = AbsoluteAlchemicalFactory(unmodified_system, ligand_atoms=alchemical_atoms, annihilate_electrostatics=True, annihilate_sterics=True, alchemical_bonds=None, alchemical_angles=None, softcore_beta=0.0)
 
         # Return the alchemically-modified system in fully-interacting form.
         alchemical_system = alchemical_factory.createPerturbedSystem()
@@ -296,25 +298,31 @@ class NCMCEngine(object):
         if direction not in ['insert', 'delete']:
             raise Exception("'direction' must be one of ['insert', 'delete']; was '%s' instead" % direction)
 
+        assert quantity_is_finite(initial_positions) == True
+
+        # Select reference topology, indices, and system based on whether we are deleting or inserting.
+        if direction == 'delete':
+            topology = topology_proposal.old_topology
+            indices = topology_proposal.unique_old_atoms
+            system = topology_proposal.old_system
+        elif direction == 'insert':
+            topology = topology_proposal.new_topology
+            indices = topology_proposal.unique_new_atoms
+            system = topology_proposal.new_system
+
+        # Handle special case of instantaneous insertion/deletion.
         if (self.nsteps == 0):
-            # Special case of instantaneous insertion/deletion.
+            # TODO: Check this is correct.
+            # TODO: Can we simplify this so there are not two code branches here?
             logP = 0.0
             final_positions = copy.deepcopy(initial_positions)
             from perses.tests.utils import compute_potential
-            if direction == 'delete':
-                potential = self.beta * compute_potential(topology_proposal.old_system, initial_positions, platform=self.platform)
-            elif direction == 'insert':
-                potential = self.beta * compute_potential(topology_proposal.new_system, initial_positions, platform=self.platform)
+            potential = self.beta * compute_potential(system, initial_positions, platform=self.platform)
             return [final_positions, logP, potential]
 
         # Create alchemical system.
         [unmodified_system, alchemical_system] = self.make_alchemical_system(topology_proposal, direction=direction)
 
-        # DEBUG: Compute initial potential of unmodified system and alchemical system to make sure finite.
-        from perses.tests.utils import compute_potential
-        #print(compute_potential(unmodified_system, initial_positions, platform=self.platform))
-        #print(compute_potential(alchemical_system, initial_positions, platform=self.platform))
-
         # Select subset of switching functions based on which alchemical parameters are present in the system.
         available_parameters = self._getAvailableParameters(alchemical_system)
         functions = { parameter_name : self.functions[parameter_name] for parameter_name in self.functions if (parameter_name in available_parameters) }
@@ -326,9 +334,11 @@ class NCMCEngine(object):
             integrator = NCMCGHMCAlchemicalIntegrator(self.temperature, alchemical_system, functions, nsteps=self.nsteps, steps_per_propagation=self.steps_per_propagation, timestep=self.timestep, direction=direction)
         else:
             raise Exception("integrator_type '%s' unknown" % self.integrator_type)
+
         # Set the constraint tolerance if specified.
         if self.constraint_tolerance is not None:
             integrator.setConstraintTolerance(self.constraint_tolerance)
+
         # Create a context on the specified platform.
         if self.platform is not None:
             context = openmm.Context(alchemical_system, integrator, self.platform)
@@ -336,78 +346,42 @@ class NCMCEngine(object):
             context = openmm.Context(alchemical_system, integrator)
         context.setPositions(initial_positions)
         context.applyConstraints(integrator.getConstraintTolerance())
+
         # Set velocities to temperature and apply velocity constraints.
         context.setVelocitiesToTemperature(self.temperature)
         context.applyVelocityConstraints(integrator.getConstraintTolerance())
 
-        # Set initial context parameters.
-        if direction == 'insert':
-            for parameter_name in available_parameters:
-                context.setParameter(parameter_name, 0)
-        elif direction == 'delete':
-            for parameter_name in available_parameters:
-                context.setParameter(parameter_name, 1)
-
-        # Compute initial potential of alchemical state.
-        initial_potential = self.beta * context.getState(getEnergy=True).getPotentialEnergy()
-        #print("Initial potential is %s" % str(initial_potential))
-        if np.isnan(initial_potential):
-            raise NaNException("Initial potential of 'insert' operation is NaN (unmodified potential was %.3f kT, alchemical potential was %.3f kT before changing lambda)" % (unmodified_potential, alchemical_potential))
-        from perses.tests.utils import compute_potential_components
-        #print("initial potential before '%s' : %f kT" % (direction, initial_potential))
-        #print("initial potential components:   %s" % str(compute_potential_components(context))) # DEBUG
-        self.write_pdb_interval = False
         # Integrate switching
         try:
-            # Write PDB file if requested.
-            if self.write_pdb_interval:
-                if direction == 'insert':
-                    topology = topology_proposal.new_topology
-                    indices = topology_proposal.unique_new_atoms
-                else:
-                    topology = topology_proposal.old_topology
-                    indices = topology_proposal.unique_old_atoms
-
-                # Write atom indices that are changing
-                import pickle
-                filename = 'ncmc-%s-%d-atomindices.pkl' % (direction, self.nattempted)
-                outfile = open(filename, 'wb')
-                pickle.dump(indices, outfile)
-                outfile.close()
-
-                from simtk.openmm.app import PDBFile
-                filename = 'ncmc-%s-%d.pdb' % (direction, self.nattempted)
-                outfile = open(filename, 'w')
-                PDBFile.writeHeader(topology, file=outfile)
-                modelIndex = 0
-                PDBFile.writeModel(topology, context.getState(getPositions=True).getPositions(asNumpy=True), file=outfile, modelIndex=modelIndex)
-                try:
-                    for step in range(self.nsteps):
-                        integrator.step(1)
-                        if (step+1)%self.write_pdb_interval == 0:
-                            modelIndex += 1
-                            # TODO: Replace with storage layer
-                            PDBFile.writeModel(topology, context.getState(getPositions=True).getPositions(asNumpy=True), file=outfile, modelIndex=modelIndex)
-                except ValueError as e:
-                    # System is exploding and coordinates won't fit in PDB ATOM fields
-                    print(e)
-
-                PDBFile.writeFooter(topology, file=outfile)
-                outfile.close()
-            else:
-                work = np.zeros([self.nsteps+1], np.float64) # work[n] is the accumulated work up to step n
-                for step in range(self.nsteps):
-                    integrator.step(1)
-                    #potential = self.beta * context.getState(getEnergy=True).getPotentialEnergy()
-                    #print("Potential at step %d is %s" % (step, str(potential)))
-                    #current_step = integrator.get_step()
-                    #print("and the integrator's current step is %d" % current_step)
+            # Write atom indices that are changing.
+            if self._storage:
+                self._storage.write_object('atomindices', indices, iteration=iteration)
+
+            # Allocate storage for work.
+            work = np.zeros([self.nsteps+1], np.float64) # work[n] is the accumulated work up to step n
+
+            # Write trajectory frame.
+            if self._storage and self.write_ncmc_interval:
+                positions = context.getState(getPositions=True).getPositions(asNumpy=True)
+                self._storage.write_configuration('positions', positions, topology, iteration=iteration, frame=0, nframes=(nsteps+1))
 
-                    # Store accumulated work
-                    work[step+1] = - integrator.getLogAcceptanceProbability(context)
+            # Perform NCMC integration.
+            for step in range(self.nsteps):
+                # Take a step.
+                integrator.step(1)
 
-                if self._storage:
-                    self._storage.write_array('work_%s' % direction, work, iteration=iteration)
+                # Store accumulated work
+                work[step+1] = integrator.getWork(context)
+
+                # Write trajectory frame.
+                if self._storage and self.write_ncmc_interval and (self.write_ncmc_interval % (step+1) == 0):
+                    positions = context.getState(getPositions=True).getPositions(asNumpy=True)
+                    assert quantity_is_finite(positions) == True
+                    self._storage.write_configuration('positions', positions, topology, iteration=iteration, frame=(step+1), nframes=(nsteps+1))
+
+            # Store work values.
+            if self._storage:
+                self._storage.write_array('work_%s' % direction, work, iteration=iteration)
 
         except Exception as e:
             # Trap NaNs as a special exception (allowing us to reject later, if desired)
@@ -421,32 +395,32 @@ class NCMCEngine(object):
                 traceback.print_exc()
                 raise e
 
-        # Set final context parameters.
-        if direction == 'insert':
-            for parameter_name in available_parameters:
-                context.setParameter(parameter_name, 1)
-        elif direction == 'delete':
-            for parameter_name in available_parameters:
-                context.setParameter(parameter_name, 0)
-
-        # Compute final potential of alchemical state.
-        final_potential = self.beta * context.getState(getEnergy=True).getPotentialEnergy()
-        if np.isnan(final_potential):
-            raise NaNException("Final potential of %s operation is NaN" % direction)
-        #print("final potential before '%s' : %f kT" % (direction, final_potential))
-        #print("final potential components: %s" % str(compute_potential_components(context))) # DEBUG
-        #print('')
-
         # Store final positions and log acceptance probability.
         final_positions = context.getState(getPositions=True).getPositions(asNumpy=True)
+        assert quantity_is_finite(final_positions) == True
         logP_NCMC = integrator.getLogAcceptanceProbability(context)
-        # DEBUG
-        #logging.debug("NCMC logP %+10.1f | initial_total_energy %+10.1f kT | final_total_energy %+10.1f kT." % (logP_NCMC, integrator.getGlobalVariableByName('initial_total_energy'), integrator.getGlobalVariableByName('final_total_energy')))
-        # Clean up NCMC switching integrator.
+
+        # Get initial and final real and alchemical potentials
+        from perses.tests.utils import compute_potential
+        initial_unmodified_potential = self.beta * compute_potential(system, initial_positions, platform=self.platform)
+        initial_alchemical_potential = self.beta * integrator.getGlobalVariableByName("Einitial") * unit.kilojoules_per_mole
+        final_unmodified_potential = self.beta * compute_potential(system, final_positions, platform=self.platform)
+        final_alchemical_potential = self.beta * context.getState(getEnergy=True).getPotentialEnergy()
+
+        # Clean up.
         del context, integrator
 
+        # Check potentials are finite
+        if np.isnan(initial_unmodified_potential) or np.isnan(initial_alchemical_potential) or np.isnan(final_unmodified_potential) or np.isnan(final_alchemical_potential):
+            msg = "A required potential of %s operation is NaN:\n" % direction
+            msg += "initial_unmodified_potential: %.3f kT\n" % initial_unmodified_potential
+            msg += "initial_alchemical_potential: %.3f kT\n" % initial_alchemical_potential
+            msg += "final_alchemical_potential: %.3f kT\n" % final_alchemical_potential
+            msg += "final_unmodified_potential: %.3f kT\n" % final_unmodified_potential
+            raise NaNException(msg)
+
         # Compute contribution from transforming real system to/from alchemical system.
-        logP_alchemical_correction = self._computeAlchemicalCorrection(unmodified_system, alchemical_system, initial_positions, final_positions, direction=direction)
+        logP_alchemical_correction = (final_unmodified_potential - final_alchemical_potential) + (initial_alchemical_potential - initial_unmodified_potential)
 
         # Compute total logP
         logP = logP_NCMC + logP_alchemical_correction
@@ -455,10 +429,11 @@ class NCMCEngine(object):
         del alchemical_system
 
         # Select whether to return initial or final potential.
+        # TODO: Do we really need to return the potential now?
         if direction == 'insert':
-            potential = initial_potential
+            potential = initial_unmodified_potential
         elif direction == 'delete':
-            potential = final_potential
+            potential = final_unmodified_potential
 
         # Keep track of statistics.
         self.nattempted += 1
@@ -523,28 +498,34 @@ class NCMCAlchemicalIntegrator(openmm.CustomIntegrator):
         """
         Reset alchemical state to initial state.
         """
+        # Set the master 'lambda' alchemical parameter to the initial state
         if self.direction == 'insert':
             self.addComputeGlobal('lambda', '0.0')
         elif self.direction == 'delete':
             self.addComputeGlobal('lambda', '1.0')
 
+        # Update all slaved alchemical parameters
         self.addUpdateAlchemicalParametersStep()
 
     def addAlchemicalPerturbationStep(self):
         """
         Add alchemical perturbation step.
         """
+        # Set the master 'lambda' alchemical parameter to the current fractional state
         if self.nsteps == 0:
+            # Toggle alchemical state
             if self.direction == 'insert':
                 self.addComputeGlobal('lambda', '1.0')
             elif self.direction == 'delete':
                 self.addComputeGlobal('lambda', '0.0')
         else:
+            # Use fractional state
             if self.direction == 'insert':
                 self.addComputeGlobal('lambda', '(step+1)/nsteps')
             elif self.direction == 'delete':
                 self.addComputeGlobal('lambda', '(nsteps - step - 1)/nsteps')
 
+        # Update all slaved alchemical parameters
         self.addUpdateAlchemicalParametersStep()
 
     def addUpdateAlchemicalParametersStep(self):
@@ -572,6 +553,7 @@ class NCMCAlchemicalIntegrator(openmm.CustomIntegrator):
         """
         self.hasStatistics = True
 
+        # TODO: This could be precomputed to save time
         self.addComputePerDof("sigma", "sqrt(kT/m)")
 
         #
@@ -587,7 +569,7 @@ class NCMCAlchemicalIntegrator(openmm.CustomIntegrator):
         self.addConstrainVelocities()
 
         self.addComputeSum("ke", "0.5*m*v*v")
-        self.addComputeGlobal("Eold", "ke + energy")
+        self.addComputeGlobal("Eold_GHMC", "ke + energy")
         self.addComputePerDof("xold", "x")
         self.addComputePerDof("vold", "v")
         self.addComputePerDof("v", "v + 0.5*dt*f/m")
@@ -597,9 +579,12 @@ class NCMCAlchemicalIntegrator(openmm.CustomIntegrator):
         self.addComputePerDof("v", "v + 0.5*dt*f/m + (x-x1)/dt")
         self.addConstrainVelocities()
         self.addComputeSum("ke", "0.5*m*v*v")
-        self.addComputeGlobal("Enew", "ke + energy")
-        self.addComputeGlobal("accept", "step(exp(-(Enew-Eold)/kT) - uniform)")
+        self.addComputeGlobal("Enew_GHMC", "ke + energy")
+        # Compute acceptance probability
+        # DEBUG: Check positions are finite
+        self.addComputeGlobal("accept", "step(exp(-(Enew_GHMC-Eold_GHMC)/kT) - uniform)")
         self.beginIfBlock("accept != 1")
+        # Reject sample, inverting velcoity
         self.addComputePerDof("x", "xold")
         self.addComputePerDof("v", "-vold")
         self.endBlock()
@@ -626,6 +611,7 @@ class NCMCAlchemicalIntegrator(openmm.CustomIntegrator):
         self.setGlobalVariableByName("total_work", 0.0)
         if (self.nsteps > 0):
             self.setGlobalVariableByName("step", 0)
+            self.setGlobalVariableByName("pstep", 0)
             if self.has_statistics:
                 self.setGlobalVariableByName("naccept", 0)
                 self.setGlobalVariableByName("ntrials", 0)
@@ -636,6 +622,11 @@ class NCMCAlchemicalIntegrator(openmm.CustomIntegrator):
         else:
             return (0,0)
 
+    def getWork(self, context):
+        """Retrieve accumulated work (in units of kT)
+        """
+        return self.getGlobalVariableByName("total_work") * unit.kilojoules_per_mole / self.kT
+
     def getLogAcceptanceProbability(self, context):
         logp_accept = -1.0*self.getGlobalVariableByName("total_work") * unit.kilojoules_per_mole / self.kT
         return logp_accept
@@ -743,6 +734,7 @@ class NCMCVVAlchemicalIntegrator(NCMCAlchemicalIntegrator):
         self.addGlobalVariable("Eold", 0)  # old energy
         self.addGlobalVariable("Enew", 0)  # new energy
         self.addGlobalVariable('kinetic', 0.0) # kinetic energy
+        self.addGlobalVariable("Einitial", 0) # initial energy after setting initial alchemical state
 
         # VV variables
         if (nsteps > 0):
@@ -760,6 +752,7 @@ class NCMCVVAlchemicalIntegrator(NCMCAlchemicalIntegrator):
 
         if nsteps == 0:
             self.addAlchemicalResetStep()
+            self.addComputeGlobal('Einitial', 'energy') # store initial energy after setting initial alchemical state
             self.addComputeGlobal("Eold", "energy")
             self.addAlchemicalPerturbationStep()
             self.addComputeGlobal("Enew", "energy")
@@ -768,6 +761,8 @@ class NCMCVVAlchemicalIntegrator(NCMCAlchemicalIntegrator):
             self.addComputeGlobal('pstep', '0')
             # Initial step only
             self.beginIfBlock('step = 0')
+            self.addAlchemicalResetStep()
+            self.addComputeGlobal('Einitial', 'energy') # store initial energy after setting initial alchemical state
             self.addComputeSum('kinetic', '0.5 * m * v^2')
             self.addComputeGlobal("Eold", "energy + kinetic")
             self.beginWhileBlock('pstep < psteps')
@@ -840,9 +835,12 @@ class NCMCGHMCAlchemicalIntegrator(NCMCAlchemicalIntegrator):
         self.addGlobalVariable('total_work', 0.0) # initial total energy (kinetic + potential)
         self.addGlobalVariable("Eold", 0)  # old energy
         self.addGlobalVariable("Enew", 0)  # new energy
+        self.addGlobalVariable("Einitial", 0) # initial energy after setting initial alchemical state
 
         if (nsteps > 0):
             # GHMC variables
+            self.addGlobalVariable("Eold_GHMC", 0)  # old GHMC energy
+            self.addGlobalVariable("Enew_GHMC", 0)  # new GHMC energy
             self.addGlobalVariable('nsteps', nsteps) # total number of NCMC steps to perform
             self.addGlobalVariable('step', 0) # current NCMC step number
             self.addPerDofVariable("x1", 0) # for velocity Verlet with constraints
@@ -865,18 +863,22 @@ class NCMCGHMCAlchemicalIntegrator(NCMCAlchemicalIntegrator):
 
         if nsteps == 0:
             self.addAlchemicalResetStep()
+            self.addComputeGlobal('Einitial', 'energy') # store initial energy after setting initial alchemical state
             self.addComputeGlobal("Eold", "energy")
             self.addAlchemicalPerturbationStep()
             self.addComputeGlobal("Enew", "energy")
             self.addComputeGlobal("total_work", "total_work + (Enew-Eold)")
         if nsteps > 0:
-            self.addComputeGlobal('pstep', '0')
+            #self.addComputeGlobal('pstep', '0')
+
             # Initial step only
             self.beginIfBlock('step = 0')
-            self.beginWhileBlock('pstep < psteps')
-            self.addGHMCStep()
-            self.addComputeGlobal('pstep', 'pstep+1')
-            self.endBlock()
+            #self.beginWhileBlock('pstep < psteps')
+            self.addAlchemicalResetStep()
+            self.addComputeGlobal('Einitial', 'energy') # store initial energy after setting initial alchemical state
+            #self.addGHMCStep()
+            #self.addComputeGlobal('pstep', 'pstep+1')
+            #self.endBlock()
             self.endBlock()
 
             # All steps
@@ -885,9 +887,9 @@ class NCMCGHMCAlchemicalIntegrator(NCMCAlchemicalIntegrator):
             self.addAlchemicalPerturbationStep()
             self.addComputeGlobal("Enew", "energy")
             self.addComputeGlobal("total_work", "total_work + (Enew-Eold)")
-            self.beginWhileBlock('pstep < psteps')
+            #self.beginWhileBlock('pstep < psteps')
             self.addGHMCStep()
-            self.addComputeGlobal('pstep', 'pstep+1')
-            self.endBlock()
+            #self.addComputeGlobal('pstep', 'pstep+1')
+            #self.endBlock()
             self.addComputeGlobal('step', 'step+1')
             self.endBlock()
diff --git a/perses/samplers/samplers.py b/perses/samplers/samplers.py
index 190ce039..9022d9b6 100644
--- a/perses/samplers/samplers.py
+++ b/perses/samplers/samplers.py
@@ -26,6 +26,7 @@ import time
 
 from perses.storage import NetCDFStorageView
 from perses.samplers import thermodynamics
+from perses.tests.utils import quantity_is_finite
 
 ################################################################################
 # LOGGER
@@ -124,6 +125,10 @@ class SamplerState(object):
 
     """
     def __init__(self, system, positions, velocities=None, box_vectors=None, platform=None):
+        assert quantity_is_finite(positions)
+        if velocities is not None:
+            assert quantity_is_finite(self.velocities)
+
         self.system = copy.deepcopy(system)
         self.positions = positions
         self.velocities = velocities
@@ -198,6 +203,9 @@ class SamplerState(object):
         self.total_energy = self.potential_energy + self.kinetic_energy
         self.volume = thermodynamics.volume(self.box_vectors)
 
+        assert quantity_is_finite(self.positions)
+        assert quantity_is_finite(self.velocities)
+
         return self
 
     def createContext(self, integrator=None, platform=None):
@@ -362,6 +370,133 @@ class SamplerState(object):
         else:
             return False
 
+################################################################################
+# GHMC INTEGRATOR (TEMPORARY UNTIL MOVED TO OPENMMTOOLS)
+################################################################################
+
+class GHMCIntegrator(openmm.CustomIntegrator):
+
+    """
+    Generalized hybrid Monte Carlo (GHMC) integrator.
+
+    """
+
+    def __init__(self, temperature=298.0 * unit.kelvin, collision_rate=91.0 / unit.picoseconds, timestep=1.0 * unit.femtoseconds):
+        """
+        Create a generalized hybrid Monte Carlo (GHMC) integrator.
+
+        Parameters
+        ----------
+        temperature : np.unit.Quantity compatible with kelvin, default: 298*unit.kelvin
+           The temperature.
+        collision_rate : np.unit.Quantity compatible with 1/picoseconds, default: 91.0/unit.picoseconds
+           The collision rate.
+        timestep : np.unit.Quantity compatible with femtoseconds, default: 1.0*unit.femtoseconds
+           The integration timestep.
+
+        Notes
+        -----
+        This integrator is equivalent to a Langevin integrator in the velocity Verlet discretization with a
+        Metrpolization step to ensure sampling from the appropriate distribution.
+
+        Additional global variables 'ntrials' and  'naccept' keep track of how many trials have been attempted and
+        accepted, respectively.
+
+        TODO
+        ----
+        * Move initialization of 'sigma' to setting the per-particle variables.
+        * Generalize to use MTS inner integrator.
+
+        Examples
+        --------
+
+        Create a GHMC integrator.
+
+        >>> temperature = 298.0 * unit.kelvin
+        >>> collision_rate = 91.0 / unit.picoseconds
+        >>> timestep = 1.0 * unit.femtoseconds
+        >>> integrator = GHMCIntegrator(temperature, collision_rate, timestep)
+
+        References
+        ----------
+        Lelievre T, Stoltz G, and Rousset M. Free Energy Computations: A Mathematical Perspective
+        http://www.amazon.com/Free-Energy-Computations-Mathematical-Perspective/dp/1848162472
+
+        """
+
+        # Initialize constants.
+        kT = kB * temperature
+        gamma = collision_rate
+
+        # Create a new custom integrator.
+        super(GHMCIntegrator, self).__init__(timestep)
+
+        #
+        # Integrator initialization.
+        #
+        self.addGlobalVariable("kT", kT)  # thermal energy
+        self.addGlobalVariable("b", np.exp(-gamma * timestep))  # velocity mixing parameter
+        self.addPerDofVariable("sigma", 0) # velocity standard deviation
+        self.addGlobalVariable("ke", 0)  # kinetic energy
+        self.addPerDofVariable("vold", 0)  # old velocities
+        self.addPerDofVariable("xold", 0)  # old positions
+        self.addGlobalVariable("Eold", 0)  # old energy
+        self.addGlobalVariable("Enew", 0)  # new energy
+        self.addGlobalVariable("accept", 0)  # accept or reject
+        self.addGlobalVariable("naccept", 0)  # number accepted
+        self.addGlobalVariable("ntrials", 0)  # number of Metropolization trials
+        self.addPerDofVariable("x1", 0)  # position before application of constraints
+
+        #
+        # Pre-computation.
+        # This only needs to be done once.
+        # TODO: Change this to setPerDofVariableByName("sigma", unit.sqrt(kT / mass).value_in_unit_system(unit.md_unit_system))
+        #
+        self.addComputePerDof("sigma", "sqrt(kT/m)")
+
+        #
+        # Velocity randomization
+        #
+        self.addComputePerDof("v", "sqrt(b)*v + sqrt(1-b)*sigma*gaussian")
+        self.addConstrainVelocities()
+
+        #
+        # Metropolized symplectic step.
+        #
+        self.addConstrainPositions()
+        self.addConstrainVelocities()
+
+        self.addComputeSum("ke", "0.5*m*v*v")
+        self.addComputeGlobal("Eold", "ke + energy")
+        self.addComputePerDof("xold", "x")
+        self.addComputePerDof("vold", "v")
+        self.addComputePerDof("v", "v + 0.5*dt*f/m")
+        self.addComputePerDof("x", "x + v*dt")
+        self.addComputePerDof("x1", "x")
+        self.addConstrainPositions()
+        self.addComputePerDof("v", "v + 0.5*dt*f/m + (x-x1)/dt")
+        self.addConstrainVelocities()
+        self.addComputeSum("ke", "0.5*m*v*v")
+        self.addComputeGlobal("Enew", "ke + energy")
+        # TODO: Check if accept/reject logic correctly handles nans
+        self.addComputeGlobal("accept", "step(exp(-(Enew-Eold)/kT) - uniform)")
+        self.beginIfBlock("accept != 1")
+        self.addComputePerDof("x", "xold")
+        self.addComputePerDof("v", "-vold")
+        self.endBlock()
+
+        #
+        # Velocity randomization
+        #
+        self.addComputePerDof("v", "sqrt(b)*v + sqrt(1-b)*sigma*gaussian")
+        self.addConstrainVelocities()
+
+        #
+        # Accumulate statistics.
+        #
+        self.addComputeGlobal("naccept", "naccept + accept")
+        self.addComputeGlobal("ntrials", "ntrials + 1")
+
 ################################################################################
 # MCMC SAMPLER
 ################################################################################
@@ -401,7 +536,7 @@ class MCMCSampler(object):
     >>> sampler.run()
 
     """
-    def __init__(self, thermodynamic_state, sampler_state, topology=None, storage=None):
+    def __init__(self, thermodynamic_state, sampler_state, topology=None, storage=None, integrator_name='GHMC'):
         """
         Create an MCMC sampler.
 
@@ -413,9 +548,10 @@ class MCMCSampler(object):
             The initial sampler state to simulate from.
         topology : simtk.openmm.app.Topology, optional, default=None
             Topology object corresponding to system being simulated (for writing)
-
         storage : NetCDFStorage, optional, default=None
             Storage layer to use for writing.
+        integrator_name : str, optional, default='GHMC'
+            Name of the integrator to use for propagation.
 
         """
         # Keep copies of initializing arguments.
@@ -423,6 +559,7 @@ class MCMCSampler(object):
         self.thermodynamic_state = copy.deepcopy(thermodynamic_state)
         self.sampler_state = copy.deepcopy(sampler_state)
         self.topology = topology
+        self.integrator_name = integrator_name
 
         self.storage = None
         if storage is not None:
@@ -435,7 +572,6 @@ class MCMCSampler(object):
         self.timestep = 1.0 * unit.femtoseconds
         self.nsteps = 500 # number of steps per update
         self.verbose = True
-        self.integrator_name = 'Langevin'
 
     def update(self):
         """
@@ -447,7 +583,8 @@ class MCMCSampler(object):
 
         # Create an integrator
         if self.integrator_name == 'GHMC':
-            from openmmtools.integrators import GHMCIntegrator
+            # TODO: Migrate GHMCIntegrator back to openmmtools
+            #from openmmtools.integrators import GHMCIntegrator
             integrator = GHMCIntegrator(temperature=self.thermodynamic_state.temperature, collision_rate=self.collision_rate, timestep=self.timestep)
             if self.verbose: print("Taking %d steps of GHMC..." % self.nsteps)
         elif self.integrator_name == 'Langevin':
@@ -698,6 +835,9 @@ class ExpandedEnsembleSampler(object):
         """
         Sample the thermodynamic state.
         """
+
+        initial_time = time.time()
+
         # Check that system and topology have same number of atoms.
         old_system = self.sampler.sampler_state.system
         old_topology = self.topology
@@ -831,10 +971,14 @@ class ExpandedEnsembleSampler(object):
                 self.nrejected += 1
                 if self.verbose: print("    rejected")
 
+            elapsed_time = time.time() - initial_time
+
             # Write to storage.
             if self.storage:
+                self.storage.write_quantity('update_state_elapsed_time', elapsed_time, iteration=self.iteration)
                 self.storage.write_configuration('positions', self.sampler.sampler_state.positions, self.sampler.topology, iteration=self.iteration)
                 self.storage.write_object('state_key', self.state_key, iteration=self.iteration)
+                self.storage.write_object('proposed_state_key', topology_proposal.new_chemical_state_key, iteration=self.iteration)
                 self.storage.write_quantity('naccepted', self.naccepted, iteration=self.iteration)
                 self.storage.write_quantity('nrejected', self.nrejected, iteration=self.iteration)
                 self.storage.write_quantity('logp_accept', logp_accept, iteration=self.iteration)
diff --git a/perses/tests/test_analysis.py b/perses/tests/test_analysis.py
index 7af8cbac..28e1d8e2 100644
--- a/perses/tests/test_analysis.py
+++ b/perses/tests/test_analysis.py
@@ -72,5 +72,5 @@ def test_analysis():
         analysis.plot_ncmc_work('ncmc.pdf')
 
 if __name__ == '__main__':
-    analysis = Analysis('output.nc')
-    analysis.plot_ncmc_work('ncmc.pdf')
+    analysis = Analysis('output-10000.nc')
+    analysis.plot_ncmc_work('ncmc-10000.pdf')
diff --git a/perses/tests/testsystems.py b/perses/tests/testsystems.py
index 89977549..9868eef3 100644
--- a/perses/tests/testsystems.py
+++ b/perses/tests/testsystems.py
@@ -1672,6 +1672,14 @@ class T4LysozymeInhibitorsTestSystem(SmallMoleculeLibraryTestSystem):
         # Intialize
         super(T4LysozymeInhibitorsTestSystem, self).__init__(**kwargs)
 
+class FusedRingsTestSystem(SmallMoleculeLibraryTestSystem):
+    """
+    Simple test system containing fused rings (benzene <--> naphtalene) in explicit solvent.
+    """
+    def __init__(self, **kwargs):
+        self.molecules = ['c1ccccc1', 'c1ccc2ccccc2c1'] # benzene, naphthalene
+        super(FusedRingsTestSystem, self).__init__(**kwargs)
+
 class ValenceSmallMoleculeLibraryTestSystem(PersesTestSystem):
     """
     Create a consistent set of samplers useful for testing SmallMoleculeProposalEngine on alkanes with a valence-only forcefield.
@@ -1958,7 +1966,7 @@ def test_valence_write_pdb_ncmc_switching():
     environment = 'vacuum'
     testsystem.exen_samplers[environment].ncmc_engine.nsteps = 10
     testsystem.exen_samplers[environment].ncmc_engine.timestep = 1.0 * unit.femtoseconds
-    testsystem.exen_samplers[environment].ncmc_engine.write_pdb_interval = 1 # write PDB files for NCMC switching
+    testsystem.exen_samplers[environment].ncmc_engine.write_ncmc_interval = 1 # write PDB files for NCMC switching
     testsystem.mcmc_samplers[environment].nsteps = 10
     testsystem.mcmc_samplers[environment].timestep = 1.0 * unit.femtoseconds
     testsystem.exen_samplers[environment].run(niterations=1)
@@ -1976,7 +1984,7 @@ def run_abl_affinity_write_pdb_ncmc_switching():
         testsystem.exen_samplers[environment].ncmc_engine.nsteps = 10000
         testsystem.exen_samplers[environment].ncmc_engine.timestep = 1.0 * unit.femtoseconds
         testsystem.exen_samplers[environment].accept_everything = False # accept everything that doesn't lead to NaN for testing
-        testsystem.exen_samplers[environment].ncmc_engine.write_pdb_interval = 100 # write PDB files for NCMC switching
+        testsystem.exen_samplers[environment].ncmc_engine.write_ncmc_interval = 100 # write PDB files for NCMC switching
         testsystem.mcmc_samplers[environment].nsteps = 10000
         testsystem.mcmc_samplers[environment].timestep = 1.0 * unit.femtoseconds
 
@@ -2011,7 +2019,7 @@ def run_constph_abl():
         testsystem.exen_samplers[environment].ncmc_engine.nsteps = 50
         testsystem.exen_samplers[environment].ncmc_engine.timestep = 1.0 * unit.femtoseconds
         testsystem.exen_samplers[environment].accept_everything = False # accept everything that doesn't lead to NaN for testing
-        #testsystem.exen_samplers[environment].ncmc_engine.write_pdb_interval = 100 # write PDB files for NCMC switching
+        #testsystem.exen_samplers[environment].ncmc_engine.write_ncmc_interval = 100 # write PDB files for NCMC switching
         testsystem.mcmc_samplers[environment].nsteps = 2500
         testsystem.mcmc_samplers[environment].timestep = 1.0 * unit.femtoseconds
 
@@ -2053,7 +2061,7 @@ def run_imidazole():
         testsystem.exen_samplers[environment].ncmc_engine.nsteps = 500
         testsystem.exen_samplers[environment].ncmc_engine.timestep = 1.0 * unit.femtoseconds
         testsystem.exen_samplers[environment].accept_everything = False # accept everything that doesn't lead to NaN for testing
-        #testsystem.exen_samplers[environment].ncmc_engine.write_pdb_interval = 100 # write PDB files for NCMC switching
+        #testsystem.exen_samplers[environment].ncmc_engine.write_ncmc_interval = 100 # write PDB files for NCMC switching
         testsystem.mcmc_samplers[environment].nsteps = 500
         testsystem.mcmc_samplers[environment].timestep = 1.0 * unit.femtoseconds
 
@@ -2066,10 +2074,38 @@ def run_imidazole():
     testsystem.sams_samplers['explicit-imidazole'].verbose=True
     testsystem.sams_samplers['explicit-imidazole'].run(niterations=100)
 
+def run_fused_rings():
+    """
+    Run fused rings test system.
+    Vary number of NCMC steps
+
+    """
+    #nsteps_to_try = [1, 10, 100, 1000, 10000, 100000] # number of NCMC steps
+    nsteps_to_try = [10, 100, 1000, 10000, 100000] # number of NCMC steps
+    for ncmc_steps in nsteps_to_try:
+        storage_filename = 'output-%d.nc' % ncmc_steps
+        testsystem = FusedRingsTestSystem(storage_filename=storage_filename)
+        for environment in ['explicit', 'vacuum']:
+            testsystem.mcmc_samplers[environment].verbose = True
+            testsystem.mcmc_samplers[environment].nsteps = 500
+            testsystem.exen_samplers[environment].verbose = True
+            testsystem.exen_samplers[environment].ncmc_engine.nsteps = ncmc_steps # NCMC switching steps
+            testsystem.exen_samplers[environment].ncmc_engine.verbose = True # verbose output of work
+            testsystem.sams_samplers[environment].verbose = True
+        testsystem.designer.verbose = True
+        testsystem.designer.run(niterations=100)
+
+        # Analyze data.
+        from perses.analysis import Analysis
+        analysis = Analysis(storage_filename=storage_filename)
+        #analysis.plot_sams_weights('sams.pdf')
+        analysis.plot_ncmc_work('ncmc-%d.pdf' % ncmc_steps)
+
 if __name__ == '__main__':
+    run_fused_rings()
     #run_valence_system()
     #run_t4_inhibitors()
-    run_imidazole()
+    #run_imidazole()
     #run_constph_imidazole()
     #run_constph_abl()
     #run_abl_affinity_write_pdb_ncmc_switching()
diff --git a/perses/tests/utils.py b/perses/tests/utils.py
index 5a69f015..ba9d2585 100644
--- a/perses/tests/utils.py
+++ b/perses/tests/utils.py
@@ -39,6 +39,25 @@ beta = 1.0/kT
 
 # TODO: Move some of these utility routines to openmoltools.
 
+def quantity_is_finite(quantity):
+    """
+    Check that elements in quantity are all finite.
+
+    Parameters
+    ----------
+    quantity : simtk.unit.Quantity
+        The quantity to check
+
+    Returns
+    -------
+    is_finite : bool
+        If quantity is finite, returns True; otherwise False.
+
+    """
+    if np.any( np.isnan( np.array(quantity / quantity.unit) ) ):
+        return False
+    return True
+
 def show_topology(topology):
     output = ""
     for atom in topology.atoms():
diff --git a/setup.py b/setup.py
index ee3cac5b..5d6c534e 100644
--- a/setup.py
+++ b/setup.py
@@ -80,14 +80,16 @@ setup(name='perses',
       zip_safe=False,
       ext_modules=extensions,
       install_requires=[
-        'openmm >=7.0',
+        'openmm >=7.0.1',
         'numpy',
         'scipy',
         'numexpr',
         'openmoltools',
-        'alchemy >=1.1.3',
+        'alchemy >=1.2.0',
+        'openmmtools',
         'numexpr',
         'netcdf4',
+        'seaborn',
         #'numba'
         ],
       )
['perses/tests/testsystems.py', 'perses/annihilation/ncmc_switching.py', 'perses/samplers/samplers.py', 'setup.py', 'perses/tests/utils.py', 'devtools/conda-recipe/meta.yaml', 'perses/analysis/analysis.py', 'perses/tests/test_analysis.py']