formatting updates to alchem tutorial

tutorials/alchemical-free-energy/files/alchemical-example.py

@@ -1,12 +1,14 @@
 from simtk import openmm, unit
 
 # Create a Lennard-Jones fluid
-pressure = 80*unit.atmospheres
-temperature = 120*unit.kelvin
-collision_rate = 5/unit.picoseconds
-timestep = 2.5*unit.femtoseconds
+pressure = 80 * unit.atmospheres
+temperature = 120 * unit.kelvin
+collision_rate = 5 / unit.picoseconds
+timestep = 2.5 * unit.femtoseconds
 from openmmtools.testsystems import LennardJonesFluid
-sigma = 3.4*unit.angstrom; epsilon = 0.238 * unit.kilocalories_per_mole
+
+sigma = 3.4 * unit.angstrom
+epsilon = 0.238 * unit.kilocalories_per_mole
 fluid = LennardJonesFluid(sigma=sigma, epsilon=epsilon)
 [system, positions] = [fluid.system, fluid.positions]
 
@@ -15,24 +17,24 @@
 system.addForce(barostat)
 
 # Retrieve the NonbondedForce
-forces = { force.__class__.__name__ : force for force in system.getForces() }
-nbforce = forces['NonbondedForce']
+forces = {force.__class__.__name__: force for force in system.getForces()}
+nbforce = forces["NonbondedForce"]
 
 # Add a CustomNonbondedForce to handle only alchemically-modified interactions
 alchemical_particles = set([0])
 chemical_particles = set(range(system.getNumParticles())) - alchemical_particles
-energy_function = 'lambda*4*epsilon*x*(x-1.0); x = (sigma/reff_sterics)^6;'
-energy_function += 'reff_sterics = sigma*(0.5*(1.0-lambda) + (r/sigma)^6)^(1/6);'
-energy_function += 'sigma = 0.5*(sigma1+sigma2); epsilon = sqrt(epsilon1*epsilon2);'
+energy_function = "lambda*4*epsilon*x*(x-1.0); x = (sigma/reff_sterics)^6;"
+energy_function += "reff_sterics = sigma*(0.5*(1.0-lambda) + (r/sigma)^6)^(1/6);"
+energy_function += "sigma = 0.5*(sigma1+sigma2); epsilon = sqrt(epsilon1*epsilon2);"
 custom_force = openmm.CustomNonbondedForce(energy_function)
-custom_force.addGlobalParameter('lambda', 1.0)
-custom_force.addPerParticleParameter('sigma')
-custom_force.addPerParticleParameter('epsilon')
+custom_force.addGlobalParameter("lambda", 1.0)
+custom_force.addPerParticleParameter("sigma")
+custom_force.addPerParticleParameter("epsilon")
 for index in range(system.getNumParticles()):
     [charge, sigma, epsilon] = nbforce.getParticleParameters(index)
     custom_force.addParticle([sigma, epsilon])
     if index in alchemical_particles:
-        nbforce.setParticleParameters(index, charge*0, sigma, epsilon*0)
+        nbforce.setParticleParameters(index, charge * 0, sigma, epsilon * 0)
 custom_force.addInteractionGroup(alchemical_particles, chemical_particles)
 system.addForce(custom_force)
 
@@ -42,43 +44,46 @@
 context.setPositions(positions)
 
 # Minimize energy
-print('Minimizing energy...')
+print("Minimizing energy...")
 openmm.LocalEnergyMinimizer.minimize(context)
 
 # Collect data
-nsteps = 2500 # number of steps per sample
-niterations = 50 # number of samples to collect per alchemical state
+nsteps = 2500  # number of steps per sample
+niterations = 50  # number of samples to collect per alchemical state
 import numpy as np
-lambdas = np.linspace(1.0, 0.0, 10) # alchemical lambda schedule
+
+lambdas = np.linspace(1.0, 0.0, 10)  # alchemical lambda schedule
 nstates = len(lambdas)
-u_kln = np.zeros([nstates,nstates,niterations], np.float64)
+u_kln = np.zeros([nstates, nstates, niterations], np.float64)
 kT = unit.AVOGADRO_CONSTANT_NA * unit.BOLTZMANN_CONSTANT_kB * integrator.getTemperature()
+
 for k in range(nstates):
     for iteration in range(niterations):
-        print('state %5d iteration %5d / %5d' % (k, iteration, niterations))
+        print(f"state {k + 1} / {nstates} iteration {iteration + 1} / {niterations}")
         # Set alchemical state
-        context.setParameter('lambda', lambdas[k])
+        context.setParameter("lambda", lambdas[k])
         # Run some dynamics
         integrator.step(nsteps)
         # Compute energies at all alchemical states
         for l in range(nstates):
-            context.setParameter('lambda', lambdas[l])
-            u_kln[k,l,iteration] = context.getState(getEnergy=True).getPotentialEnergy() / kT
+            context.setParameter("lambda", lambdas[l])
+            u_kln[k, l, iteration] = context.getState(getEnergy=True).getPotentialEnergy() / kT
 
 # Estimate free energy of Lennard-Jones particle insertion
 from pymbar import MBAR, timeseries
+
 # Subsample data to extract uncorrelated equilibrium timeseries
-N_k = np.zeros([nstates], np.int32) # number of uncorrelated samples
+N_k = np.zeros([nstates], np.int32)  # number of uncorrelated samples
 for k in range(nstates):
-    [nequil, g, Neff_max] = timeseries.detectEquilibration(u_kln[k,k,:])
-    indices = timeseries.subsampleCorrelatedData(u_kln[k,k,:], g=g)
+    [nequil, g, Neff_max] = timeseries.detectEquilibration(u_kln[k, k, :])
+    indices = timeseries.subsampleCorrelatedData(u_kln[k, k, :], g=g)
     N_k[k] = len(indices)
-    u_kln[k,:,0:N_k[k]] = u_kln[k,:,indices].T
+    u_kln[k, :, 0 : N_k[k]] = u_kln[k, :, indices].T
 # Compute free energy differences and statistical uncertainties
 mbar = MBAR(u_kln, N_k)
 [DeltaF_ij, dDeltaF_ij, Theta_ij] = mbar.getFreeEnergyDifferences()
 
-print('DeltaF_ij (kT):')
+print("DeltaF_ij (kT):")
 print(DeltaF_ij)
-print('dDeltaF_ij (kT):')
+print("dDeltaF_ij (kT):")
 print(dDeltaF_ij)