Moving scripts of benchmarking to pymbar-datasets

scripts/benchmark_covariance.py

@@ -0,0 +1,45 @@
+import pandas as pd
+import numpy as np
+import pymbar
+from pymbar_datasets import load_gas_data, load_8proteins_data
+import time
+
+def load_oscillators(n_states, n_samples):
+    name = "%dx%d oscillators" % (n_states, n_samples)
+    O_k = np.linspace(1, 5, n_states)
+    k_k = np.linspace(1, 3, n_states)
+    N_k = (np.ones(n_states) * n_samples).astype('int')
+    test = pymbar.testsystems.harmonic_oscillators.HarmonicOscillatorsTestCase(O_k, k_k)
+    x_n, u_kn, N_k_output, s_n = test.sample(N_k, mode='u_kn')
+    return name, u_kn, N_k_output, s_n
+
+def load_exponentials(n_states, n_samples):
+    name = "%dx%d exponentials" % (n_states, n_samples)
+    rates = np.linspace(1, 3, n_states)
+    N_k = (np.ones(n_states) * n_samples).astype('int')
+    test = pymbar.testsystems.exponential_distributions.ExponentialTestCase(rates)
+    x_n, u_kn, N_k_output, s_n = test.sample(N_k, mode='u_kn')
+    return name, u_kn, N_k_output, s_n
+
+
+mbar_gens = {"new":lambda u_kn, N_k: pymbar.MBAR(u_kn, N_k)}
+systems = [lambda : load_exponentials(25, 100), lambda : load_exponentials(100, 100), lambda : load_exponentials(250, 250),
+lambda : load_oscillators(25, 100), lambda : load_oscillators(100, 100), lambda : load_oscillators(250, 250),
+lambda : load_oscillators(500, 100), lambda : load_oscillators(1000, 50), lambda : load_oscillators(2000, 20), lambda : load_oscillators(4000, 10), 
+lambda : load_exponentials(500, 100), lambda : load_exponentials(1000, 50), lambda : load_exponentials(2000, 20), lambda : load_oscillators(4000, 10), 
+load_gas_data, load_8proteins_data]
+
+timedata = []
+for version, mbar_gen in mbar_gens.items():
+    for sysgen in systems:
+        name, u_kn, N_k, s_n = sysgen()
+        K, N = u_kn.shape
+        mbar = mbar_gen(u_kn, N_k)
+        time0 = time.time()
+        fij, dfij = mbar.getFreeEnergyDifferences(uncertainty_method="svd-ew-kab")
+        dt =  time.time() - time0
+        timedata.append([name, K, N, dt])
+
+
+timedata = pd.DataFrame(timedata, columns=["name", "K", "N", "time"])
+print timedata.to_string(float_format=lambda x: "%.3g" % x)

scripts/benchmark_mbar.py

@@ -0,0 +1,46 @@
+import pandas as pd
+import numpy as np
+import pymbar
+from pymbar_datasets import load_gas_data, load_8proteins_data, load_k69_data
+import time
+
+def load_oscillators(n_states, n_samples):
+    name = "%dx%d oscillators" % (n_states, n_samples)
+    O_k = np.linspace(1, 5, n_states)
+    k_k = np.linspace(1, 3, n_states)
+    N_k = (np.ones(n_states) * n_samples).astype('int')
+    test = pymbar.testsystems.harmonic_oscillators.HarmonicOscillatorsTestCase(O_k, k_k)
+    x_n, u_kn, N_k_output, s_n = test.sample(N_k, mode='u_kn')
+    return name, u_kn, N_k_output, s_n
+
+def load_exponentials(n_states, n_samples):
+    name = "%dx%d exponentials" % (n_states, n_samples)
+    rates = np.linspace(1, 3, n_states)
+    N_k = (np.ones(n_states) * n_samples).astype('int')
+    test = pymbar.testsystems.exponential_distributions.ExponentialTestCase(rates)
+    x_n, u_kn, N_k_output, s_n = test.sample(N_k, mode='u_kn')
+    return name, u_kn, N_k_output, s_n
+
+
+solver_protocol = None
+mbar_gens = {"new":lambda u_kn, N_k, x_kindices: pymbar.MBAR(u_kn, N_k, x_kindices=x_kindices)}
+systems = [lambda : load_exponentials(25, 100), lambda : load_exponentials(100, 100), lambda : load_exponentials(250, 250),
+lambda : load_oscillators(25, 100), lambda : load_oscillators(100, 100), lambda : load_oscillators(250, 250),
+load_gas_data, load_8proteins_data, load_k69_data]
+
+timedata = []
+for version, mbar_gen in mbar_gens.items():
+    for sysgen in systems:
+        name, u_kn, N_k, s_n = sysgen()
+        time0 = time.time()
+        mbar = mbar_gen(u_kn, N_k, s_n)
+        dt = time.time() - time0
+        wsum = np.linalg.norm(np.exp(mbar.Log_W_nk).sum(0) - 1.0)
+        wdot = np.linalg.norm(np.exp(mbar.Log_W_nk).dot(N_k) - 1.0)
+        obj, grad = pymbar.mbar_solvers.mbar_objective_and_gradient(u_kn, N_k, mbar.f_k)
+        grad_norm = np.linalg.norm(grad)
+        timedata.append([name, version, dt, grad_norm, wsum, wdot])
+
+
+timedata = pd.DataFrame(timedata, columns=["name", "version", "time", "|grad|", "|W.sum(0) - 1|", "|W.dot(N_k) - 1|"])
+print timedata.to_string(float_format=lambda x: "%.3g" % x)

scripts/time_expectations.py

@@ -0,0 +1,33 @@
+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("u_kn.npz")["arr_0"]
+N_k = np.load("N_k.npz")["arr_0"]
+s_n = np.load("s_n.npz")["arr_0"]
+
+mbar = pymbar.MBAR(u_kn, N_k, s_n)
+mbar0 = pymbar.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)