#!/usr/local/bin/env python

#=============================================================================================
# MODULE DOCSTRING
#=============================================================================================

"""
Estimate steady-state nonequilibrium free energy of a TIP3P water box as a function of number of water molecules and timestep.

Use mpi4py to parallelize water box sizes.

"""

#=============================================================================================
# GLOBAL IMPORTS
#=============================================================================================

import sys
import math
import doctest
import time
import numpy

import simtk.unit as units
import simtk.openmm as openmm

import netCDF4 as netcdf 

#=============================================================================================
# CONSTANTS
#=============================================================================================

kB = units.BOLTZMANN_CONSTANT_kB * units.AVOGADRO_CONSTANT_NA

#=============================================================================================
# INTEGRATORS
#=============================================================================================

def GHMCIntegrator(timestep, temperature=298.0*units.kelvin, gamma=50.0/units.picoseconds):
    """
    Create a generalized hybrid Monte Carlo (GHMC) integrator.
    
    ARGUMENTS

    timestep (numpy.unit.Quantity compatible with femtoseconds) - the integration timestep
    temperature (numpy.unit.Quantity compatible with kelvin) - the temperature
    gamma (numpy.unit.Quantity compatible with 1/picoseconds) - the collision rate

    RETURNS

    integrator (simtk.openmm.CustomIntegrator) - a GHMC integrator

    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.

    """

    kT = kB * temperature
        
    integrator = openmm.CustomIntegrator(timestep)

    integrator.addGlobalVariable("kT", kB*temperature) # thermal energy
    integrator.addGlobalVariable("b", numpy.exp(-gamma*timestep)) # velocity mixing parameter
    integrator.addPerDofVariable("sigma", 0) 
    integrator.addGlobalVariable("ke", 0) # kinetic energy
    integrator.addPerDofVariable("vold", 0) # old velocities
    integrator.addPerDofVariable("xold", 0) # old positions
    integrator.addGlobalVariable("Eold", 0) # old energy
    integrator.addGlobalVariable("Enew", 0) # new energy
    integrator.addGlobalVariable("accept", 0) # accept or reject
    integrator.addGlobalVariable("naccept", 0) # number accepted
    integrator.addGlobalVariable("ntrials", 0) # number of Metropolization trials
    integrator.addPerDofVariable("x1", 0) # position before application of constraints
    
    #
    # Pre-computation.
    # This only needs to be done once, but it needs to be done for each degree of freedom.
    # Could move this to initialization?
    #
    integrator.addComputePerDof("sigma", "sqrt(kT/m)")

    # 
    # Velocity perturbation.
    #
    integrator.addComputePerDof("v", "sqrt(b)*v + sqrt(1-b)*sigma*gaussian")
    integrator.addConstrainVelocities();
    
    #
    # Metropolized symplectic step.
    #
    integrator.addUpdateContextState();
    integrator.addComputeSum("ke", "0.5*m*v*v")
    integrator.addComputeGlobal("Eold", "ke + energy")
    integrator.addComputePerDof("xold", "x")
    integrator.addComputePerDof("vold", "v")
    integrator.addComputePerDof("v", "v + 0.5*dt*f/m")
    integrator.addComputePerDof("x", "x + v*dt")
    integrator.addComputePerDof("x1", "x")
    integrator.addConstrainPositions();
    integrator.addComputePerDof("v", "v + 0.5*dt*f/m + (x-x1)/dt")
    integrator.addConstrainVelocities();
    integrator.addComputeSum("ke", "0.5*m*v*v")
    integrator.addComputeGlobal("Enew", "ke + energy")
    integrator.addComputeGlobal("accept", "step(exp(-(Enew-Eold)/kT) - uniform)")
    integrator.addComputePerDof("x", "x*accept + xold*(1-accept)")
    integrator.addComputePerDof("v", "v*accept - vold*(1-accept)")

    #
    # Velocity randomization
    #
    integrator.addComputePerDof("v", "sqrt(b)*v + sqrt(1-b)*sigma*gaussian")
    integrator.addConstrainVelocities();

    #
    # Accumulate statistics.
    #
    integrator.addComputeGlobal("naccept", "naccept + accept")
    integrator.addComputeGlobal("ntrials", "ntrials + 1")   

    return integrator

def VVVRIntegrator(timestep, temperature=298.0*units.kelvin, gamma=50.0/units.picoseconds, timestep_correction=True):
    """
    Create a velocity verlet with velocity randomization (VVVR) integrator.
    
    ARGUMENTS

    timestep (numpy.unit.Quantity compatible with femtoseconds) - the integration timestep
    temperature (numpy.unit.Quantity compatible with kelvin) - the temperature
    gamma (numpy.unit.Quantity compatible with 1/picoseconds) - the collision rate

    RETURNS

    integrator (simtk.openmm.CustomIntegrator) - a VVVR integrator

    NOTES
    
    This integrator is equivalent to a Langevin integrator in the velocity Verlet discretization with a
    timestep correction to ensure that the field-free diffusion constant is timestep invariant.

    The global 'pseudowork' keeps track of the pseudowork accumulated during integration, and can be
    used to correct the sampled statistics or in a Metropolization scheme.
    
    TODO

    Move initialization of 'sigma' to setting the per-particle variables.
    We can ditch pseudowork and instead use total energy difference - heat.
    
    """

    kT = kB * temperature
    print "kT = %.3f kcal/mol" % (kT / units.kilocalories_per_mole) # DEBUG
    
    a = numpy.exp(-gamma*timestep)

    b = 1.0
    if timestep_correction:
        b = (2.0/(gamma*timestep)) * numpy.tanh(gamma*timestep/2.0)
    print "Using VVVR constants a = %f | b = %f" % (a, b)

    integrator = openmm.CustomIntegrator(timestep)
    
    integrator.addGlobalVariable("kT", kT) # thermal energy
    integrator.addPerDofVariable("sigma", 0) 
    integrator.addPerDofVariable("x1", 0) # position before application of constraints
    integrator.addGlobalVariable("a", a) # timestep_correction
    integrator.addGlobalVariable("b", b) # timestep_correction

    #                                     
    # Pre-computation.
    #
    integrator.addComputePerDof("sigma", "sqrt(kT/m)")

    # 
    # Velocity perturbation.
    #
    integrator.addComputePerDof("v", "sqrt(a)*v + sqrt(1-a)*sigma*gaussian")
    integrator.addConstrainVelocities();
    
    #
    # Symplectic step with constraints.
    #
    integrator.addUpdateContextState(); 
    integrator.addComputePerDof("v", "v + 0.5*(b*dt)*f/m")
    integrator.addComputePerDof("x", "x + (b*dt)*v")
    integrator.addComputePerDof("x1", "x")
    integrator.addConstrainPositions();
    integrator.addComputePerDof("v", "v + 0.5*(b*dt)*f/m + (x-x1)/(b*dt)")
    integrator.addConstrainVelocities();

    #
    # Velocity randomization
    #
    integrator.addComputePerDof("v", "sqrt(a)*v + sqrt(1-a)*sigma*gaussian")
    integrator.addConstrainVelocities();

    return integrator
    
#=============================================================================================
# UTILITY SUBROUTINES
#=============================================================================================

def initialize_netcdf(ncfile, system, timestep, temperature, gamma, nsteps):
    """
    Initialize NetCDF file for storage.
    
    """    
    
    # Create dimensions.
    ncfile.createDimension('samples', 0) # unlimited number of samples
    ncfile.createDimension('atoms', system.getNumParticles()) # number of particles in system
    ncfile.createDimension('dimensions', 3) # number of spatial dimensions
    ncfile.createDimension('timesteps', len(timesteps_to_try)) # number of timesteps to try
    ncfile.createDimension('single', 1)

    # Create variables.
    ncvar_positions = ncfile.createVariable('positions', 'f', ('samples', 'atoms', 'dimensions'))
    ncvar_velocities = ncfile.createVariable('velocities', 'f', ('samples', 'atoms', 'dimensions'))
    ncvar_box_vectors = ncfile.createVariable('box_vectors', 'f', ('samples','dimensions','dimensions'))        
    ncvar_volumes  = ncfile.createVariable('volumes', 'f', ('samples',))
    ncvar_potential_energy = ncfile.createVariable('potential_energy', 'f', ('samples',))        
    ncvar_kinetic_energy = ncfile.createVariable('kinetic_energy', 'f', ('samples',))        
    ncvar_simulation_time = ncfile.createVariable('simulation_time', 'f', ('samples',))        
    ncvar_nsteps = ncfile.createVariable('nsteps', 'i', ('samples',))        

    # Serialize OpenMM System object.
    ncvar_serialized_state = ncfile.createVariable('system', str, ('single',), zlib=True)
    ncvar_serialized_state[0] = system.__getstate__()
    
    # Define units for variables.
    setattr(ncvar_positions, 'units', 'nm')
    setattr(ncvar_velocities, 'units', 'nm/ps')
    setattr(ncvar_box_vectors, 'units', 'nm')
    setattr(ncvar_volumes, 'units', 'nm**3')
    setattr(ncvar_potential_energy, 'units', 'kT')
    setattr(ncvar_kinetic_energy, 'units', 'kT')
    setattr(ncvar_simulation_time, 'units', 'ps')

    # Define long (human-readable) names for variables.
    setattr(ncvar_positions, "long_name", "positions[sample,particle,dimenson] is position of coordinate 'dimension' of particle 'particle' for sample 'sample'.")
    setattr(ncvar_velocities, "long_name", "velocities[sample,particle,dimension] is velocity of coordinate 'dimension' of particle 'particle' for sample 'sample.")
    setattr(ncvar_box_vectors, "long_name", "box_vectors[sample,i,j] is dimension j of box vector i for sample 'sample'.")
    setattr(ncvar_volumes, "long_name", "volume[sample] is the box volume for sample 'sample'.")
    setattr(ncvar_simulation_time, "long_name", "simulation_time[sample] is the simulation time for sample 'sample'.")
    setattr(ncvar_nsteps, "long_name", "nsteps is the number of steps between writes.")
    
    # Create timestamp variable.
    ncvar_timestamp = ncfile.createVariable('timestamp', str, ('samples',))

    # Store timestep and gamma.
    ncvar_timestep = ncfile.createVariable('timestep', 'f', ('single',))
    ncvar_timestep[0] = timestep / units.femtoseconds
    setattr(ncvar_timestep, 'units', 'fs')
    ncvar_gamma = ncfile.createVariable('gamma', 'f', ('single',))
    ncvar_gamma[0] = gamma * units.picoseconds
    setattr(ncvar_gamma, 'units', '1/ps')
    ncvar_nsteps[0] = nsteps
    setattr(ncvar_nsteps, 'units', 'none')
    ncvar_temperature = ncfile.createVariable('temperature', 'f', ('single',))
    ncvar_temperature[0] = temperature / units.kelvin
    setattr(ncvar_temperature, 'units', 'kelvin')
    
    # Force sync to disk to avoid data loss.
    ncfile.sync()
        
    return

#=============================================================================================
# PARAMETERS
#=============================================================================================

# Sets of parameters to regress over.
timestep_correction_flags_to_try = [True, False]
systems_to_try = ['IdealGas', 'LennardJonesFluid', 'WaterBox'] 
timesteps_to_try = units.Quantity([0.5, 1.0, 2.0, 4.0, 8.0, 16.0], units.femtoseconds) # MD timesteps
gammas_to_try = units.Quantity([0.01, 0.1, 1.0, 10.0, 100.0, 1000.0], units.picoseconds**-1) # collision rates

ngpus = 4

print "timestep corrections to try: %s" % str(timestep_correction_flags_to_try)
print "systems to try: %s" % str(systems_to_try)
print "timesteps to try: %s" % str(timesteps_to_try)
print "gammas to try: %s" % str(gammas_to_try)

ghmc_nsteps = 10000 # number of steps to generate new uncorrelated sample with GHMC
ghmc_timestep = 0.5 * units.femtoseconds
write_interval = 1.0 * units.picoseconds # time between writing data
nsamples = 2000 # number of samples to generate
nequil = 100 # number of NPT equilibration iterations

# DEBUG
#systems_to_try = ['LennardJonesFluid']
#systems_to_try = ['IdealGas']
#gammas_to_try = units.Quantity([10.0], units.picoseconds**-1) # collision rates
#timestep_correction_flags_to_try = [True]
#timesteps_to_try = units.Quantity([1.0], units.femtoseconds) # MD timesteps
#nsamples = 101 # number of samples to generate
#nequil = 10 # number of NPT equilibration iterations

verbose = True
verbose = False

#=============================================================================================
# Initialize MPI.
#=============================================================================================

try:
    from mpi4py import MPI # MPI wrapper
    rank = MPI.COMM_WORLD.rank
    size = MPI.COMM_WORLD.size
    print "Node %d / %d" % (MPI.COMM_WORLD.rank, MPI.COMM_WORLD.size)
except:
    print "mpi4py not available---using serial execution."
    rank = 0
    size = 1
    
print "%d / %d" % (rank, size)

platform_name = 'CUDA'
#platform_name = 'CPU'
#platform_name = 'OpenCL'
platform = openmm.Platform.getPlatformByName(platform_name)
deviceid = rank % ngpus
#deviceid = 2 # DEBUG
if platform_name == 'CUDA':
    platform.setPropertyDefaultValue('CudaDeviceIndex', '%d' % deviceid) # select Cuda device index
if platform_name == 'OpenCL':
    platform.setPropertyDefaultValue('OpenCLDeviceIndex', '%d' % deviceid) # select OpenCL device index

print "node %3d using GPU %d" % (rank, deviceid)

noptionsets = len(timestep_correction_flags_to_try) * len(systems_to_try) * len(timesteps_to_try) * len(gammas_to_try)
if rank == 0: print "There are %d option sets to try." % noptionsets
for index in range(rank, noptionsets, size):

    #=============================================================================================
    # Select problem to work on.
    #=============================================================================================

    def select_options(options_list, index):
        selected_options = list()
        for option in options_list:
            noptions = len(option)
            selected_options.append(option[index % noptions])
            index = int(index/noptions)
        return selected_options

    try:
        [use_timestep_correction, system_name, timestep, gamma] = select_options([timestep_correction_flags_to_try, systems_to_try, timesteps_to_try, gammas_to_try], index)
    except:
        continue

    # Compute number of steps per write.
    nsteps = int(write_interval / timestep)

    # Create filename to store data in.
    store_filename = 'data/vvvr-%s-%s-%s-%s.nc' % (str(use_timestep_correction), system_name, str(timestep/units.femtoseconds), str(gamma*units.picoseconds))
    print "Node %d: store_filename = %s" % (rank, store_filename)

    #=============================================================================================
    # Create system to simulate.
    #=============================================================================================

    import systems
    constructor = getattr(systems, system_name)
    if system_name == 'LennardJonesFluid':
        # Lennard-Jones fluid parameters (argon).
        mass     = 39.9 * units.amu
        sigma    = 3.4 * units.angstrom
        epsilon  = 120.0 * units.kelvin * kB

        # Set temperature, pressure, and collision rate for stochastic thermostats.
        temperature = 0.9 / (kB / epsilon)
        pressure = 4.0 / (sigma**3 / epsilon) / units.AVOGADRO_CONSTANT_NA

        [system, positions] = constructor(nx=6, ny=6, nz=6, mass=mass, sigma=sigma, epsilon=epsilon, switch=False, dispersion_correction=True)
    else:
        temperature = 298.0 * units.kelvin
        pressure = 1.0 * units.atmosphere # pressure for equilibration
        [system, positions] = constructor()

    kT = kB * temperature # thermal energy
    beta = 1.0 / kT # inverse temperature

    ndof = 3*system.getNumParticles() - system.getNumConstraints()
    nparticles = system.getNumParticles()

    print "Node %d: Box has %d particles" % (rank, nparticles)

    # Turn off output from all but head node.
    if rank != 0: verbose = False

    #=============================================================================================
    # Open NetCDF file for writing.
    #=============================================================================================

    import os.path
    resume = False
    if os.path.exists(store_filename):
        print "Node %d: Attempting to resume from file '%s'..." % (rank, store_filename)
        ncfile = netcdf.Dataset(store_filename, 'a')
        resume = True
    else:
        print "Node %d: Opening '%s' for writing..." % (rank, store_filename)
        ncfile = netcdf.Dataset(store_filename, 'w', version='NETCDF4')
        initialize_netcdf(ncfile, system, timestep, temperature, gamma, nsteps)    

    #=============================================================================================
    # Select precision.
    #=============================================================================================
    
    options = {}
    if platform_name == "OpenCL":
        options = {"OpenCLPrecision" : "double"}
    if platform_name == "CUDA":
        options = {"CudaPrecision" : "double"}
    tolerance = 1.0e-8 # constraint tolerance

    #=============================================================================================
    # Equilibrate with GHMC using a barostat.
    #=============================================================================================

    if not resume:
        # Equilibrate with Monte Carlo barostat.
        import copy
        system_with_barostat = copy.deepcopy(system)
        barostat = openmm.MonteCarloBarostat(pressure, temperature)
        system_with_barostat.addForce(barostat)
        ghmc_integrator = GHMCIntegrator(ghmc_timestep, temperature=temperature)
        ghmc_integrator.setConstraintTolerance(tolerance)
        ghmc_global_variables = { ghmc_integrator.getGlobalVariableName(index) : index for index in range(ghmc_integrator.getNumGlobalVariables()) }
        ghmc_context = openmm.Context(system_with_barostat, ghmc_integrator, platform, options)
        
        # Modify random number seeds to be unique.
        seed = ghmc_integrator.getRandomNumberSeed()
        print seed
        ghmc_integrator.setRandomNumberSeed(seed + rank)
        barostat.setRandomNumberSeed(seed + rank + size)
            
        # Set positions and velocities.
        ghmc_context.setPositions(positions)

        # Minimize.
        print "node %3d minimizing..." % rank
        openmm.LocalEnergyMinimizer.minimize(ghmc_context, 0.01 * units.kilocalories_per_mole / units.nanometer, 20)
        print "node %3d minimization complete." % rank

        # Compute initial volume.
        state = ghmc_context.getState()
        box_vectors = state.getPeriodicBoxVectors(asNumpy=True)
        volume = box_vectors[0,0] * box_vectors[1,1] * box_vectors[2,2]
        print "node %d: initial volume %8.3f nm^3" % (rank, volume / units.nanometers**3)

        # Equilibrate system with NPT.
        volume_history = numpy.zeros([nequil], numpy.float64)
        for iteration in range(nequil):
            ghmc_integrator.setGlobalVariable(ghmc_global_variables['naccept'], 0)
            ghmc_integrator.setGlobalVariable(ghmc_global_variables['ntrials'], 0)

            # Generate new sample from equilibrium distribution with GHMC.
            ghmc_integrator.step(ghmc_nsteps)
    
            # Compute volume.
            state = ghmc_context.getState(getEnergy=True)
            box_vectors = state.getPeriodicBoxVectors(asNumpy=True)
            potential = state.getPotentialEnergy()
            kinetic = state.getKineticEnergy()
            volume = box_vectors[0,0] * box_vectors[1,1] * box_vectors[2,2]
            volume_history[iteration] = volume / units.nanometers**3
            max_radius = box_vectors[0,0] / 2.0 # half the box width
            instantaneous_kinetic_temperature = kinetic / (1.5 * system.getNumParticles() * kB)

            naccept = ghmc_integrator.getGlobalVariable(ghmc_global_variables['naccept'])
            ntrials = ghmc_integrator.getGlobalVariable(ghmc_global_variables['ntrials'])
            fraction_accepted = float(naccept) / float(ntrials)
            print "%64s: GHMC equil %5d / %5d | accepted %6d / %6d (%7.3f %%) | volume %8.3f nm^3 | max radius %8.3f nm | potential %8.3f kT | temperature %8.3f K" % (store_filename, iteration, nequil, naccept, ntrials, fraction_accepted*100.0, volume / units.nanometers**3, max_radius / units.nanometers, potential / kT, instantaneous_kinetic_temperature / units.kelvin)
        
        # Extract coordinates and box vectors, along with final energy.
        state = ghmc_context.getState(getPositions=True, getVelocities=True, getEnergy=True)
        positions = state.getPositions(asNumpy=True)
        velocities = state.getVelocities(asNumpy=True)
        box_vectors = state.getPeriodicBoxVectors(asNumpy=True)
        volume = box_vectors[0,0] * box_vectors[1,1] * box_vectors[2,2]
        final_potential = state.getPotentialEnergy()

        del ghmc_context, ghmc_integrator

    #=============================================================================================
    # Create and destry dummy context.
    #=============================================================================================
    
    #print "node %d: creating dummy context..." % rank
    #integrator = openmm.VerletIntegrator(timestep)
    #context = openmm.Context(system, integrator, platform)
    #del context, integrator

    #=============================================================================================
    # Resume if NetCDF file already exists.
    #=============================================================================================

    first_sample = 0 # sample index to start from
    simulation_time = 0.0 * units.picoseconds # simulation time
    if resume:
        # Get positions and velocities from NetCDF file.
        sample = ncfile.variables['positions'].shape[0] - 1 # back up one sample, in case it was incomplete
        if (sample > 0):
            print "Resuming from NetCDF file at sample %d" % sample
            positions = units.Quantity(ncfile.variables['positions'][sample,:,:], units.nanometers)
            velocities = units.Quantity(ncfile.variables['velocities'][sample,:,:], units.nanometers / units.picoseconds)
            box_vectors = units.Quantity(ncfile.variables['box_vectors'][sample,:,:], units.nanometers)
            simulation_time = units.Quantity(float(ncfile.variables['simulation_time'][sample]), units.picoseconds)
            nequil = 0 # don't equilibrate
            first_sample = sample # start at 'sample'
            final_potential = None

    #=============================================================================================
    # Run production simulation.
    #=============================================================================================

    print "node %3d starting production simulation..." % rank

    # Initialize VVVR integrator and context.
    # TODO: Initialize number of steps and simulation time.
    integrator = VVVRIntegrator(timestep, temperature=temperature, gamma=gamma, timestep_correction=use_timestep_correction)
    integrator.setConstraintTolerance(tolerance)
    context = openmm.Context(system, integrator, platform, options)
    context.setPeriodicBoxVectors(box_vectors[0,:], box_vectors[1,:], box_vectors[2,:])  # this must come first!
    context.setPositions(positions)
    context.setVelocities(velocities) 
    context.setTime(simulation_time)

    # Modify random number seeds to be unique.
    seed = integrator.getRandomNumberSeed()
    print seed
    integrator.setRandomNumberSeed(seed + rank)

    print "node %3d context created" % rank

    # DEBUG
    #print box_vectors
    # Extract coordinates and box vectors.
    #state = context.getState(getPositions=True, getVelocities=True, getEnergy=True)
    #current_positions = state.getPositions(asNumpy=True)
    #box_vectors = state.getPeriodicBoxVectors(asNumpy=True)
    #potential_energy = state.getPotentialEnergy()
    #kinetic_energy = state.getKineticEnergy()
    #volume = box_vectors[0,0] * box_vectors[1,1] * box_vectors[2,2]
    #simulation_time = state.getTime()
    #print box_vectors
    #elapsed_time = 0.0
    #print "%64s: sample %8d/%8d | potential %12.3f kT | kinetic %12.3f kT | total %12.3f K | %12.3f s elapsed this iteration" % (store_filename, -1, nsamples, potential_energy/kT, kinetic_energy/kT, instantaneous_kinetic_temperature / units.kelvin, elapsed_time)

    # Consistency check
    #if final_potential:
    #    if abs((final_potential - potential_energy) / kT) > 0.01:
    #        print "***********************************************"
    #        print "Significant difference in potential energy!"
    #        print "old positions"
    #        print positions
    #        print "current positions"
    #        print current_positions
    #        print "***********************************************"

    for sample in range(first_sample, nsamples):
        initial_time = time.time()
        
        # Run VVVR.
        integrator.step(nsteps)

        # Extract coordinates and box vectors.
        state = context.getState(getPositions=True, getVelocities=True, getEnergy=True)
        positions = state.getPositions(asNumpy=True)
        velocities = state.getVelocities(asNumpy=True)
        box_vectors = state.getPeriodicBoxVectors(asNumpy=True)
        potential_energy = state.getPotentialEnergy()
        kinetic_energy = state.getKineticEnergy()
        volume = box_vectors[0,0] * box_vectors[1,1] * box_vectors[2,2]
        simulation_time = state.getTime()

        # Store data.
        ncfile.variables['positions'][sample,:,:] = positions[:,:] / units.nanometers
        ncfile.variables['velocities'][sample,:,:] = velocities[:,:] / (units.nanometers/units.picoseconds)
        ncfile.variables['box_vectors'][sample,:,:] = box_vectors[:,:] / units.nanometers
        ncfile.variables['volumes'][sample] = volume / (units.nanometers**3)
        ncfile.variables['potential_energy'][sample] = potential_energy / kT
        ncfile.variables['kinetic_energy'][sample] = kinetic_energy / kT
        ncfile.variables['simulation_time'][sample] = simulation_time / units.picoseconds
        ncfile.variables['timestamp'][sample] = time.ctime()
        ncfile.sync()

        instantaneous_kinetic_temperature = kinetic_energy / (1.5 * system.getNumParticles() * kB)

        final_time = time.time()
        elapsed_time = final_time - initial_time
        if sample % 100 == 0: print "%64s: sample %8d/%8d | potential %12.3f kT | kinetic %12.3f kT | total %12.3f K | %12.3f s elapsed this iteration" % (store_filename, sample, nsamples, potential_energy/kT, kinetic_energy/kT, instantaneous_kinetic_temperature / units.kelvin, elapsed_time)

    # Close NetCDF file.
    ncfile.close()

    # Clean up
    del context, integrator