MBAR benchmark with lots of empty states

bench.py

import pandas as pd
import numpy as np
import pymbar

samples_per_state = 250
n_states = 100
rates = np.linspace(1, 3, n_states)
N_k = np.ones(n_states, 'int') * samples_per_state
n_samples = n_states * samples_per_state

test = pymbar.testsystems.exponential_distributions.ExponentialTestCase(rates)
x_n, u_kn0, N_k_output = test.sample(N_k, mode='u_kn')

n_empty = 10000
N_k = np.pad(N_k, (0, n_empty), mode='constant')

u_kn = np.zeros((n_states + n_empty, n_samples))
u_kn[0:n_states] = u_kn0


%time mbar = pymbar.MBAR(u_kn, N_k)
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)

%timeit f2 = pymbar.mbar_solvers.self_consistent_update(u_kn, N_k, mbar.f_k)