Merge pull request #148 from kyleabeauchamp/returnvals

Eliminated variable length returns

examples/harmonic-oscillators-distributions.py

@@ -132,7 +132,7 @@
 
   # Initialize the MBAR class, determining the free energies.
   mbar = MBAR(u_kln, N_k, relative_tolerance=1.0e-10,verbose=False) # use fast Newton-Raphson solver
-  (Deltaf_ij_estimated, dDeltaf_ij_estimated) = mbar.getFreeEnergyDifferences()
+  (Deltaf_ij_estimated, dDeltaf_ij_estimated, _theta) = mbar.getFreeEnergyDifferences()
 
   # Compute error from analytical free energy differences.
   Deltaf_ij_error = Deltaf_ij_estimated - Deltaf_ij_analytical

examples/harmonic-oscillators.py

@@ -140,7 +140,7 @@ def GetAnalytical(beta,K,O,observables):
 print "      Testing getFreeEnergyDifferences       "
 print "============================================="
 
-(Delta_f_ij_estimated, dDelta_f_ij_estimated) = mbar.getFreeEnergyDifferences()
+(Delta_f_ij_estimated, dDelta_f_ij_estimated, _Theta_ij) = mbar.getFreeEnergyDifferences()
 
 # Compute error from analytical free energy differences.
 Delta_f_ij_error = Delta_f_ij_estimated - Delta_f_ij_analytical
@@ -512,7 +512,8 @@ def GetAnalytical(beta,K,O,observables):
 print "      Testing computeOverlap   "
 print "============================================"
 
-O_ij = mbar.computeOverlap(output='matrix')                
+O, O_i, O_ij = mbar.computeOverlap()
+
 print "Overlap matrix output"
 print O_ij
 
@@ -522,11 +523,11 @@ def GetAnalytical(beta,K,O,observables):
     print "looks like it is."
   else:
     print "but it's not."
-O_i = mbar.computeOverlap(output='eigenvalues')                
+
 print "Overlap eigenvalue output"
 print O_i
 
-O = mbar.computeOverlap(output='scalar')                
+
 print "Overlap scalar output"
 print O
 

pymbar/mbar.py

@@ -272,7 +272,7 @@ def __init__(self, u_kn, N_k, maximum_iterations=10000, relative_tolerance=1.0e-
 
         if self.verbose:
             print "MBAR initialization complete."
-        return
+
 
     @property
     def W_nk(self):
@@ -362,21 +362,19 @@ def computeEffectiveSampleNumber(self, verbose = False):
         return N_eff
 
     #=========================================================================
-    def computeOverlap(self, output='scalar'):
+    def computeOverlap(self):
         """Compute estimate of overlap matrix between the states.
 
         Returns
         -------
+        overlap_scalar : np.ndarray, float, shape=(K, K)
+            One minus the largest nontrival eigenvalue
+        eigenval : np.ndarray, float, shape=(K)
+            The sorted (descending) eigenvalues of the overlap matrix.
         O : np.ndarray, float, shape=(K, K)
             estimated state overlap matrix: O[i,j] is an estimate
             of the probability of observing a sample from state i in state j
 
-        Parameters
-        ----------
-        output : string, optional
-            One of 'scalar', 'matrix', 'eigenvalues', 'all', specifying
-        what measure of overlap to return
-
         Notes
         -----
 
@@ -401,14 +399,8 @@ def computeOverlap(self, output='scalar'):
         # sort in descending order
         eigenval = np.sort(eigenval)[::-1]
         overlap_scalar = 1 - eigenval[1]
-        if (output == 'scalar'):
-            return overlap_scalar
-        elif (output == 'eigenvalues'):
-            return eigenval
-        elif (output == 'matrix'):
-            return O
-        elif (output == 'all'):
-            return overlap_scalar, eigenval, O
+
+        return overlap_scalar, eigenval, O
 
     #=========================================================================
     def getFreeEnergyDifferences(self, compute_uncertainty=True, uncertainty_method=None, warning_cutoff=1.0e-10, return_theta=False):
@@ -431,11 +423,15 @@ def getFreeEnergyDifferences(self, compute_uncertainty=True, uncertainty_method=
 
         Returns
         -------
-        Deltaf_ij :L np.ndarray, float, shape=(K, K)
+        Deltaf_ij : np.ndarray, float, shape=(K, K)
             Deltaf_ij[i,j] is the estimated free energy difference
-        dDeltaf_ij :L np.ndarray, float, shape=(K, K)
+        dDeltaf_ij : np.ndarray, float, shape=(K, K)
+            If compute_uncertainty==True, 
             dDeltaf_ij[i,j] is the estimated statistical uncertainty 
-            (one standard deviation) in Deltaf_ij[i,j]
+            (one standard deviation) in Deltaf_ij[i,j].  Otherwise None.
+        Theta_ij : np.ndarray, float, shape=(K, K)
+            The theta_matrix if return_theta==True, otherwise None.
+            
 
         Notes
         -----
@@ -454,33 +450,33 @@ def getFreeEnergyDifferences(self, compute_uncertainty=True, uncertainty_method=
         >>> from pymbar import testsystems
         >>> (x_n, u_kn, N_k, s_n) = testsystems.HarmonicOscillatorsTestCase().sample(mode='u_kn')
         >>> mbar = MBAR(u_kn, N_k)
-        >>> [Deltaf_ij, dDeltaf_ij] = mbar.getFreeEnergyDifferences()
+        >>> Deltaf_ij, dDeltaf_ij, Theta_ij = mbar.getFreeEnergyDifferences()
 
         """
-
+        Deltaf_ij, dDeltaf_ij, Theta_ij = None, None, None  # By default, returns None for dDelta and Theta
+
         # Compute free energy differences.
         f_i = np.matrix(self.f_k)
         Deltaf_ij = f_i - f_i.transpose()
 
         # zero out numerical error for thermodynamically identical states
         self._zerosamestates(Deltaf_ij)
-
-        returns = []
-        returns.append(np.array(Deltaf_ij))
+
+        Deltaf_ij = np.array(Deltaf_ij)  # Convert from np.matrix to np.array
 
         if compute_uncertainty or return_theta:
             # Compute asymptotic covariance matrix.
             Theta_ij = self._computeAsymptoticCovarianceMatrix(
                 np.exp(self.Log_W_nk), self.N_k, method=uncertainty_method)
 
         if compute_uncertainty:
-            dDeltaf_ij = self._ErrorOfDifferences(Theta_ij,warning_cutoff=warning_cutoff)
+            dDeltaf_ij = self._ErrorOfDifferences(Theta_ij, warning_cutoff=warning_cutoff)
             # zero out numerical error for thermodynamically identical states
             self._zerosamestates(dDeltaf_ij)
             # Return matrix of free energy differences and uncertainties.
-            returns.append(np.array(dDeltaf_ij))
+            dDeltaf_ij = np.array(dDeltaf_ij)
 
-        return returns
+        return Deltaf_ij, dDeltaf_ij, Theta_ij
 
     #=========================================================================
     def computeExpectationsInner(self, A_n, u_ln, state_map,
@@ -828,6 +824,7 @@ def computeExpectations(self, A_n, u_kn=None, output='averages', state_dependent
             dA_i  (K np float64 array) - dA_i[i] is uncertainty estimate (one standard deviation) for A_i[i]
             or
             dA_ij (K np float64 array) - dA_ij[i,j] is uncertainty estimate (one standard deviation) for the difference in A beteen i and j
+            or None, if compute_uncertainty is False.
 
         References
         ----------
@@ -895,7 +892,8 @@ def computeExpectations(self, A_n, u_kn=None, output='averages', state_dependent
                                                       return_theta=compute_uncertainty,
                                                       uncertainty_method=uncertainty_method,
                                                       warning_cutoff=warning_cutoff)
-        returns = []
+
+        mu, sigma = None, None
 
         if compute_uncertainty:
             # we want the theta matrix for the exponentials of the
@@ -909,20 +907,19 @@ def computeExpectations(self, A_n, u_kn=None, output='averages', state_dependent
             covA_ij = np.array(Theta[0:K,0:K]+Theta[K:2*K,K:2*K]-Theta[0:K,K:2*K]-Theta[K:2*K,0:K])
 
         if output == 'averages':
-            # Return expectations and uncertainties.
-            returns.append(inner_results['observables'])
+            mu = inner_results['observables']
             if compute_uncertainty:
-                returns.append(np.sqrt(covA_ij[0:K,0:K].diagonal()))
+                sigma = np.sqrt(covA_ij[0:K,0:K].diagonal())
 
         if output == 'differences':
             A_im = np.matrix(inner_results['observables'])
             A_ij = A_im - A_im.transpose()
 
-            returns.append(np.array(A_ij))
+            mu = np.array(A_ij)
             if compute_uncertainty:
-                returns.append(self._ErrorOfDifferences(covA_ij,warning_cutoff=warning_cutoff))
+                sigma = self._ErrorOfDifferences(covA_ij,warning_cutoff=warning_cutoff)
 
-        return returns
+        return mu, sigma
 
     #=========================================================================
     def computeMultipleExpectations(self, A_in, u_n, compute_uncertainty=True, compute_covariance=False,
@@ -942,12 +939,15 @@ def computeMultipleExpectations(self, A_in, u_n, compute_uncertainty=True, compu
             A_in[i,n] = A_i(x_n), the value of phase observable i for configuration n at state of interest
         u_n : np.ndarray, float, shape=(N)
             u_n[n] is the reduced potential of configuration n at the state of interest
+        compute_uncertainty : bool, optional, default=True
+            If True, calculate the uncertainty
+        compute_covariance : bool, optional, default=False
+            If True, calculate the covariance
         uncertainty_method : string, optional
             Choice of method used to compute asymptotic covariance method, or None to use default
             See help for computeAsymptoticCovarianceMatrix() for more information on various methods. (default: None)
         warning_cutoff : float, optional
             Warn if squared-uncertainty is negative and larger in magnitude than this number (default: 1.0e-10)
-        return_covariance: Whether to return the covariance of the I observables.    
 
         Returns
         -------
@@ -956,8 +956,10 @@ def computeMultipleExpectations(self, A_in, u_n, compute_uncertainty=True, compu
             A_i[i] is the estimate for the expectation of A_i(x) at the state specified by u_kn
         dA_i : np.ndarray, float, shape = (I)
             dA_i[i] is the uncertainty in the expectation of A_state_map[i](x) at the state specified by u_n[state_map[i],:] 
+            or None if compute_uncertainty is False
         d2A_ij : np.ndarray, float, shape=(I, I)
             d2A_ij[i,j] is the COVARIANCE in the estimates of A_i[i] and A_i[j]: we can't actually take a square root
+            or None if compute_covariance is False
 
         Examples
         --------
@@ -967,7 +969,7 @@ def computeMultipleExpectations(self, A_in, u_n, compute_uncertainty=True, compu
         >>> mbar = MBAR(u_kn, N_k)
         >>> A_in = np.array([x_n,x_n**2,x_n**3])
         >>> u_n = u_kn[0,:]
-        >>> [A_i, d2A_ij] = mbar.computeMultipleExpectations(A_in, u_kn)
+        >>> (A_i, dA_i, d2A_ij) = mbar.computeMultipleExpectations(A_in, u_kn)
 
         """
 
@@ -992,8 +994,9 @@ def computeMultipleExpectations(self, A_in, u_n, compute_uncertainty=True, compu
                                                       return_theta=(compute_uncertainty or compute_covariance),
                                                       uncertainty_method=uncertainty_method,
                                                       warning_cutoff=warning_cutoff)
-        returns = []
-        returns.append(inner_results['observables'])
+
+        expectations, uncertainties, covariances = None, None, None
+        expectations = inner_results['observables']
 
         if compute_uncertainty or compute_covariance:
             Adiag = np.zeros([2*I,2*I],dtype=np.float64)
@@ -1004,14 +1007,14 @@ def computeMultipleExpectations(self, A_in, u_n, compute_uncertainty=True, compu
 
         if compute_uncertainty:
             covA_ij = np.array(Theta[0:I,0:I]+Theta[I:2*I,I:2*I]-Theta[0:I,I:2*I]-Theta[I:2*I,0:I])
-            returns.append(np.sqrt(covA_ij[0:I,0:I].diagonal()))
+            uncertainties = np.sqrt(covA_ij[0:I,0:I].diagonal())
 
         if compute_covariance:
             # compute estimate of statistical covariance of the observables
-            returns.append(inner_results['Theta'][0:I,0:I])
+            covariances = inner_results['Theta'][0:I,0:I]
 
-        # Return expectations, uncertainties, covariances, theta
-        return returns
+        return expectations, uncertainties, covariances
+
 
     #=========================================================================
     def computePerturbedFreeEnergies(self, u_ln, compute_uncertainty=True, uncertainty_method=None, warning_cutoff=1.0e-10):
@@ -1024,7 +1027,7 @@ def computePerturbedFreeEnergies(self, u_ln, compute_uncertainty=True, uncertain
         u_ln : np.ndarray, float, shape=(L, Nmax)
             u_ln[l,n] is the reduced potential energy of uncorrelated
             configuration n evaluated at new state k.  Can be completely indepednent of the original number of states.
-        compute_uncertainty : bool, optional
+        compute_uncertainty : bool, optional, default=True
             If False, the uncertainties will not be computed (default: True)
         uncertainty_method : string, optional
             Choice of method used to compute asymptotic covariance method, or None to use default
@@ -1038,13 +1041,14 @@ def computePerturbedFreeEnergies(self, u_ln, compute_uncertainty=True, uncertain
             Deltaf_ij[i,j] = f_j - f_i, the dimensionless free energy difference between new states i and j
         dDeltaf_ij : np.ndarray, float, shape=(L, L)
             dDeltaf_ij[i,j] is the estimated statistical uncertainty in Deltaf_ij[i,j]
+            or None if compute_uncertainty is False
 
         Examples
         --------
         >>> from pymbar import testsystems
         >>> (x_n, u_kn, N_k, s_n) = testsystems.HarmonicOscillatorsTestCase().sample(mode='u_kn')
         >>> mbar = MBAR(u_kn, N_k)
-        >>> [Deltaf_ij, dDeltaf_ij] = mbar.computePerturbedFreeEnergies(u_kn)
+        >>> Deltaf_ij, dDeltaf_ij = mbar.computePerturbedFreeEnergies(u_kn)
         """
 
         # Convert to np matrix.
@@ -1067,16 +1071,16 @@ def computePerturbedFreeEnergies(self, u_ln, compute_uncertainty=True, uncertain
                                                       uncertainty_method=uncertainty_method,
                                                       warning_cutoff=warning_cutoff)
 
-        returns = []
+        Deltaf_ij, dDeltaf_ij = None, None
+
         f_k = np.matrix(inner_results['free energies'])
         Deltaf_ij = np.array(f_k - f_k.transpose())
-        returns.append(Deltaf_ij)
 
         if (compute_uncertainty):
-            returns.append(self._ErrorOfDifferences(inner_results['Theta'],warning_cutoff=warning_cutoff))
+            dDeltaf_ij = self._ErrorOfDifferences(inner_results['Theta'],warning_cutoff=warning_cutoff)
 
         # Return matrix of free energy differences and uncertainties.
-        return returns
+        return Deltaf_ij, dDeltaf_ij
 
     #=====================================================================
 
@@ -1370,7 +1374,6 @@ def computePMF(self, u_n, bin_n, nbins, uncertainties='from-lowest', pmf_referen
         else:
             raise "Uncertainty method '%s' not recognized." % uncertainties
 
-        return
 
     #=========================================================================
     # PRIVATE METHODS - INTERFACES ARE NOT EXPORTED

pymbar/tests/test_covariance.py

@@ -30,8 +30,8 @@ def _test(data_generator):
         raise(SkipTest("Error importing pytables to load dataset; skipping test. Error was '%s'" % exc))
     print(name)
     mbar = pymbar.MBAR(U, N_k)
-    fij1, dfij1 = mbar.getFreeEnergyDifferences(uncertainty_method="svd")
-    fij2, dfij2 = mbar.getFreeEnergyDifferences(uncertainty_method="svd-ew")
+    fij1, dfij1, Theta_ij = mbar.getFreeEnergyDifferences(uncertainty_method="svd")
+    fij2, dfij2, Theta_ij = mbar.getFreeEnergyDifferences(uncertainty_method="svd-ew")
     eq(pymbar.mbar_solvers.mbar_gradient(U, N_k, mbar.f_k), np.zeros(N_k.shape), decimal=8)
     eq(np.exp(mbar.Log_W_nk).sum(0), np.ones(len(N_k)), decimal=10)
     eq(np.exp(mbar.Log_W_nk).dot(N_k), np.ones(U.shape[1]), decimal=10)

pymbar/tests/test_mbar.py

@@ -71,7 +71,7 @@ def test_mbar_free_energies():
         eq(N_k, N_k_output)
         mbar = MBAR(u_kn, N_k)
 
-        fe, fe_sigma = mbar.getFreeEnergyDifferences()
+        fe, fe_sigma, Theta_ij = mbar.getFreeEnergyDifferences()
         fe, fe_sigma = fe[0,1:], fe_sigma[0,1:]
 
         fe0 = test.analytical_free_energies()
@@ -153,7 +153,7 @@ def test_mbar_computeMultipleExpectations():
         A[0,:] = x_n
         A[1,:] = x_n**2
         state = 1
-        mu, sigma = mbar.computeMultipleExpectations(A,u_kn[state,:])
+        mu, sigma, covariances = mbar.computeMultipleExpectations(A,u_kn[state,:])
         mu0 = test.analytical_observable(observable = 'position')[state]
         mu1 = test.analytical_observable(observable = 'position^2')[state]
         z = (mu0 - mu[0]) / sigma[0]
@@ -214,33 +214,28 @@ def test_mbar_computeOverlap():
     x_n, u_kn, N_k_output, s_n = test.sample(even_N_k, mode='u_kn')
     mbar = MBAR(u_kn, even_N_k)
 
-    O = mbar.computeOverlap(output='matrix')
-    test = np.matrix((1.0/d)*np.ones([d,d]))
-    eq(O, test, decimal=precision)
+    overlap_scalar, eigenval, O = mbar.computeOverlap()
 
-    O = mbar.computeOverlap(output='eigenvalues')
-    test = np.zeros(d)
-    test[0] = 1.0
-    eq(O, test, decimal=precision)
+    reference_matrix = np.matrix((1.0/d)*np.ones([d,d]))
+    reference_eigenvalues = np.zeros(d)
+    reference_eigenvalues[0] = 1.0
+    reference_scalar = np.float64(1.0)
 
-    O = mbar.computeOverlap(output='scalar')
-    test = np.float64(1.0)
-    eq(O, test, decimal=precision)
+    eq(O, reference_matrix, decimal=precision)
+    eq(eigenval, reference_eigenvalues, decimal=precision)
+    eq(overlap_scalar, reference_scalar, decimal=precision)
 
     # test of more straightforward examples
     for system_generator in system_generators:
         name, test = system_generator()
         x_n, u_kn, N_k_output, s_n = test.sample(N_k, mode='u_kn')
         mbar = MBAR(u_kn, N_k)
-        O = mbar.computeOverlap(output='matrix')
+        overlap_scalar, eigenval, O = mbar.computeOverlap()
 
         # rows of matrix should sum to one
         sumrows = np.array(np.sum(O,axis=1))
         eq(sumrows, np.ones(np.shape(sumrows)), decimal=precision)
-
-        # first eigenvalue should be 1
-        O = mbar.computeOverlap(output='eigenvalues')
-        eq(O[0], np.float64(1.0), decimal=precision)
+        eq(eigenval[0], np.float64(1.0), decimal=precision)
 
 def test_mbar_getWeights():