Merge pull request #171 from kyleabeauchamp/fastexpectation

[WIP] Faster Expectation Calculation

pymbar/mbar.py

@@ -620,13 +620,17 @@ def computeExpectationsInner(self, A_n, u_ln, state_map,
         N_k[0:K] = self.N_k
         f_k[0:K] = self.f_k
 
+        # Pre-calculate the log denominator: Eqns 13, 14 in MBAR paper
+        states_with_samples = (self.N_k > 0)
+        log_denominator_n = logsumexp(self.f_k[states_with_samples] - self.u_kn[states_with_samples].T, b=self.N_k[states_with_samples], axis=1)
         # Compute row of W_nk matrix for the extra states corresponding to u_ln 
         # that the state list specifies
         for l in L_list:
             la = K+l  #l, augmented
-            Log_W_nk[:, la] = self._computeUnnormalizedLogWeights(u_ln[l,:]) 
-            f_k[la] = -logsumexp(Log_W_nk[:, la])
-            Log_W_nk[:, la] += f_k[la]
+            # Calculate log normalizing constants and log weights via Eqns 13, 14
+            log_C_a = -logsumexp(-u_ln[l] - log_denominator_n)
+            Log_W_nk[:, la] = log_C_a - u_ln[l] - log_denominator_n
+            f_k[la] = log_C_a
 
         # Compute the remaining rows/columns of W_nk, and calculate
         # their normalizing constants c_k

pymbar/testsystems/exponential_distributions.py

@@ -156,3 +156,49 @@ def sample(self, N_k=[10, 20, 30, 40, 50], mode='u_kln'):
             raise Exception("Unknown mode '%s'" % mode)
 
         return
+
+    @classmethod
+    def evenly_spaced_exponentials(cls, n_states, n_samples_per_state, lower_rate=1.0, upper_rate=3.0):
+        """Generate samples from evenly spaced exponential distributions.
+
+        Parameters
+        ----------
+        n_states : np.ndarray, int
+            number of states
+        n_samples_per_state : np.ndarray, int
+            number of samples per state.  The total number of samples
+            n_samples will be equal to n_states * n_samples_per_state
+        lower_O_k : float, optional, default=1.0
+            Lower bound of O_k values
+        upper_O_k : float, optional, default=5.0
+            Upper bound of O_k values
+        lower_k_k : float, optional, default=1.0
+            Lower bound of O_k values
+        upper_k_k : float, optional, default=3.0
+            Upper bound of k_k values
+
+        Returns
+        -------
+        name: str
+            Name of testsystem
+        testsystem : TestSystem
+            The testsystem object
+        x_n : np.ndarray, shape=(n_samples)
+            Coordinates of the samples
+        u_kn : np.ndarray, shape=(n_states, n_samples)
+            Reduced potential energies
+        N_k : np.ndarray, shape=(n_states)
+            Number of samples drawn from each state
+        s_n : np.ndarray, shape=(n_samples)
+            State of origin of each sample
+        """
+
+        name = "%dx%d exponentials" % (n_states, n_samples_per_state)
+
+        rates = np.linspace(lower_rate, upper_rate, n_states)
+        N_k = (np.ones(n_states) * n_samples_per_state).astype('int')
+
+        testsystem = cls(rates)
+        x_n, u_kn, N_k_output, s_n = testsystem.sample(N_k, mode='u_kn')
+
+        return name, testsystem, x_n, u_kn, N_k_output, s_n

pymbar/testsystems/harmonic_oscillators.py

@@ -164,3 +164,49 @@ def sample(self, N_k=[10, 20, 30, 40, 50], mode='u_kn'):
             raise Exception("Unknown mode '%s'" % mode)
 
         return
+
+    @classmethod
+    def evenly_spaced_oscillators(cls, n_states, n_samples_per_state, lower_O_k=1.0, upper_O_k=5.0, lower_k_k=1.0, upper_k_k=3.0):
+        """Generate samples from evenly spaced harmonic oscillators.
+
+        Parameters
+        ----------
+        n_states : np.ndarray, int
+            number of states
+        n_samples_per_state : np.ndarray, int
+            number of samples per state.  The total number of samples
+            n_samples will be equal to n_states * n_samples_per_state
+        lower_O_k : float, optional, default=1.0
+            Lower bound of O_k values
+        upper_O_k : float, optional, default=5.0
+            Upper bound of O_k values
+        lower_k_k : float, optional, default=1.0
+            Lower bound of O_k values
+        upper_k_k : float, optional, default=3.0
+            Upper bound of k_k values
+
+        Returns
+        -------
+        name: str
+            Name of testsystem
+        testsystem : TestSystem
+            The testsystem object
+        x_n : np.ndarray, shape=(n_samples)
+            Coordinates of the samples
+        u_kn : np.ndarray, shape=(n_states, n_samples)
+            Reduced potential energies
+        N_k : np.ndarray, shape=(n_states)
+            Number of samples drawn from each state
+        s_n : np.ndarray, shape=(n_samples)
+            State of origin of each sample
+        """
+        name = "%dx%d oscillators" % (n_states, n_samples_per_state)
+
+        O_k = np.linspace(lower_O_k, upper_O_k, n_states)
+        k_k = np.linspace(lower_k_k, upper_k_k, n_states)
+        N_k = (np.ones(n_states) * n_samples_per_state).astype('int')
+
+        testsystem = cls(O_k, k_k)
+        x_n, u_kn, N_k_output, s_n = testsystem.sample(N_k, mode='u_kn')
+
+        return name, testsystem, x_n, u_kn, N_k_output, s_n

scripts/misc/time_expectations.py

@@ -0,0 +1,36 @@
+import numpy as np
+import pymbar
+import time
+
+n_states, n_samples = 400, 10
+name, testsystem, x_n, u_kn, N_k, s_n = pymbar.testsystems.HarmonicOscillatorsTestCase.evenly_spaced_oscillators(n_states, n_samples)
+
+# Saving the results in case you want to use them in an old version of pymbar that doesn't have the helper function to generate data.
+np.savez("u_kn.npz", u_kn)
+np.savez("N_k.npz", N_k)
+np.savez("s_n.npz", s_n)
+
+u_kn = np.load("/home/kyleb/src/kyleabeauchamp/pymbar/u_kn.npz")["arr_0"]
+N_k = np.load("/home/kyleb/src/kyleabeauchamp/pymbar/N_k.npz")["arr_0"]
+s_n = np.load("/home/kyleb/src/kyleabeauchamp/pymbar/s_n.npz")["arr_0"]
+
+mbar = pymbar.MBAR(u_kn, N_k, s_n)
+mbar0 = pymbar.old_mbar.MBAR(u_kn, N_k, initial_f_k=mbar.f_k)  # use initial guess to speed this up
+
+N = u_kn.shape[1]
+x = np.random.normal(size=(N))
+x0 = np.random.normal(size=(N))
+x1 = np.random.normal(size=(N))
+x2 = np.random.normal(size=(N))
+
+%time fe, fe_sigma = mbar.computePerturbedFreeEnergies(np.array([x0, x1, x2]))
+%time fe0, fe_sigma0 = mbar0.computePerturbedFreeEnergies(np.array([x0, x1, x2]))
+
+fe - fe0
+fe_sigma - fe_sigma0
+
+%time A, dA = mbar.computeExpectations(x, compute_uncertainty=False)
+%time A = mbar0.computeExpectations(x, compute_uncertainty=False)
+
+
+