"""\
Package for estimating populations of discrete states based on Boltzmann occupancy.

@author John D. Chodera <jchodera@berkeley.edu>
"""

import numpy

def boltzmannPopulation(u_n):
   """Compute population of discrete Boltzmann-weighted states.

   ARGUMENTS
     u_n [numpy N-array] - u_n[n] is the (unitless) reduced potential energy of state n

   RETURNS
     p_n [numpy N-array] - p_n[n] is he fractional population of state n

   NOTES
     The reduced potential energy u_n[n] is given in terms of the inverse temperature beta
     (in units of 1/energy) multiplied by the potential or free energy of state n (in units
     of energy).  See, for example, ref. [1].

     The maximum argument is subtracted to avoid overflow in numpy.exp.

   REFERENCES
     [1] Shirts MR and Chodera JD. Statistically optimal analysis of samples from multiple equilibrium states.
     J. Chem. Phys. 129:124105, 2008
     http://dx.doi.org/10.1063/1.2978177     

   EXAMPLE
     >>> F_n = numpy.array([-2.4, -3.4, -2.2, -2.1, -1.5]) # free energies of states (kcal/mol)
     >>> kB = 1.381e-23 * 6.0221e23 / 4184.0   # Boltzmann constant (kcal/mol/K)
     >>> kT = 298.0 * kB                       # thermal energy at 298 K (kcal/mol)
     >>> beta = 1.0 / kT                       # compute inverse temperature beta (1/(kcal/mol))
     >>> u_n = beta * F_n                      # compute (unitless) reduced potential
     >>> p_n = boltzmannPopulation(u_n)        # compute Boltzmann populations
     >>> for n in range(len(p_n)): print "%3d %8.5f" % (n, p_n[n]) # print populations
       0  0.11344
       1  0.02097
       2  0.15900
       3  0.18824
       4  0.51836
   """

   # Compute maximum argument to exp to avoid overflow.
   maxarg = u_n.max()
   
   # Compute unnormalized populations.
   p_n = numpy.exp(u_n - maxarg)

   # Normalize populations.
   p_n = p_n / p_n.sum()

   # Return populations
   return p_n

#=============================================================================================
# MAIN AND TESTS
#=============================================================================================

if __name__ == "__main__":
   import doctest
   doctest.testmod()