def __init__(self, beta, K, O):
"""Generate a set of harmonic oscillators.
Parameters
----------
beta : np.ndarray, float
number of samples per state
K : np.ndarray, float
force constants of harmonic oscillators in dimensionless units
O : np.ndarray, float
offsets of the harmonic oscillators in dimensionless units
"""
self.n_states = len(K)
K = ensure_type(K, np.float64, 1, 'K')
O = ensure_type(O, np.float64, 1, 'O', self.n_states)
self.sigma = (beta * K) ** (-0.5)
self.f_k = -np.log(np.sqrt(2.0 * np.pi) * self.sigma)
f_as_2D = np.array([self.f_k])
self.f_ij = f_as_2D - f_as_2D.T
# Calculate and store observables
self.RMS_displacement = self.sigma
self.potential_energy = 1.0 / (2.0 * beta) * np.ones(self.n_states)
self.position = O.copy()
self.position_squared = (1.0 + beta * K * O ** 2.0) / (beta * K)
self.displacement = np.zeros(self.n_states) # CHECK
self.displacement_squared = self.sigma ** 2.0 # CHECK
def validate_inputs(u_kn, N_k, f_k):
"""Check types and return inputs for MBAR calculations.
Parameters
----------
u_kn or q_kn : np.ndarray, shape=(n_states, n_samples), dtype='float'
The reduced potential energies or unnormalized probabilities
N_k : np.ndarray, shape=(n_states), dtype='int'
The number of samples in each state
f_k : np.ndarray, shape=(n_states), dtype='float'
The reduced free energies of each state
Returns
-------
u_kn or q_kn : np.ndarray, shape=(n_states, n_samples), dtype='float'
The reduced potential energies or unnormalized probabilities
N_k : np.ndarray, shape=(n_states), dtype='float'
The number of samples in each state. Converted to float because this cast is required when log is calculated.
f_k : np.ndarray, shape=(n_states), dtype='float'
The reduced free energies of each state
"""
n_states, n_samples = u_kn.shape
u_kn = ensure_type(u_kn,
'float',
2,
"u_kn or Q_kn",
shape=(n_states, n_samples))
N_k = ensure_type(
N_k, 'float', 1, "N_k", shape=(n_states, ), warn_on_cast=False
) # Autocast to float because will be eventually used in float calculations.
f_k = ensure_type(f_k, 'float', 1, "f_k", shape=(n_states, ))
return u_kn, N_k, f_k
def validate_inputs(u_kn, N_k, f_k):
"""Check types and return inputs for MBAR calculations.
Parameters
----------
u_kn or q_kn : np.ndarray, shape=(n_states, n_samples), dtype='float'
The reduced potential energies or unnormalized probabilities
N_k : np.ndarray, shape=(n_states), dtype='int'
The number of samples in each state
f_k : np.ndarray, shape=(n_states), dtype='float'
The reduced free energies of each state
Returns
-------
u_kn or q_kn : np.ndarray, shape=(n_states, n_samples), dtype='float'
The reduced potential energies or unnormalized probabilities
N_k : np.ndarray, shape=(n_states), dtype='float'
The number of samples in each state. Converted to float because this cast is required when log is calculated.
f_k : np.ndarray, shape=(n_states), dtype='float'
The reduced free energies of each state
"""
n_states, n_samples = u_kn.shape
u_kn = ensure_type(u_kn, "float", 2, "u_kn or Q_kn", shape=(n_states, n_samples))
N_k = ensure_type(
N_k, "float", 1, "N_k", shape=(n_states,), warn_on_cast=False
) # Autocast to float because will be eventually used in float calculations.
f_k = ensure_type(f_k, "float", 1, "f_k", shape=(n_states,))
return u_kn, N_k, f_k
def __init__(self, beta, K, O):
"""Generate a set of harmonic oscillators.
Parameters
----------
beta : np.ndarray, float
number of samples per state
K : np.ndarray, float
force constants of harmonic oscillators in dimensionless units
O : np.ndarray, float
offsets of the harmonic oscillators in dimensionless units
"""
self.n_states = len(K)
K = ensure_type(K, np.float64, 1, 'K')
O = ensure_type(O, np.float64, 1, 'O', self.n_states)
self.sigma = (beta * K)**(-0.5)
self.f_k = -np.log(np.sqrt(2.0 * np.pi) * self.sigma)
f_as_2D = np.array([self.f_k])
self.f_ij = f_as_2D - f_as_2D.T
# Calculate and store observables
self.RMS_displacement = self.sigma
self.potential_energy = 1.0 / (2.0 * beta) * np.ones(self.n_states)
self.position = O.copy()
self.position_squared = (1.0 + beta * K * O**2.0) / (beta * K)
self.displacement = np.zeros(self.n_states) # CHECK
self.displacement_squared = self.sigma**2.0 # CHECK
def save(name, u_kn, N_k, s_n=None, least_significant_digit=None):
"""Create an HDF5 dump of an existing MBAR job for later use / testing.
Parameters
----------
name : str
Name of dataset
u_kn : np.ndarray, dtype='float', shape=(n_states, n_samples)
Reduced potential energies
N_k : np.ndarray, dtype='int', shape=(n_states)
Number of samples taken from each state
s_n : np.ndarray, optional, default=None, dtype=int, shape=(n_samples)
The state of origin of each state. If none, guess the state origins.
least_significant_digit : int, optional, default=None
If not None, perform lossy compression using tables.Filter(least_significant_digit=least_significant_digit)
Notes
-----
The output HDF5 files should be readible by the helper funtions pymbar_datasets.py
"""
import tables
(n_states, n_samples) = u_kn.shape
u_kn = ensure_type(u_kn,
'float',
2,
"u_kn or Q_kn",
shape=(n_states, n_samples))
N_k = ensure_type(N_k, 'int64', 1, "N_k", shape=(n_states, ))
if s_n is None:
s_n = get_sn(N_k)
s_n = ensure_type(s_n, 'int64', 1, "s_n", shape=(n_samples, ))
hdf_filename = os.path.join("./", "%s.h5" % name)
f = tables.File(hdf_filename, 'a')
f.createCArray("/",
"u_kn",
tables.Float64Atom(),
obj=u_kn,
filters=tables.Filters(
complevel=9,
complib="zlib",
least_significant_digit=least_significant_digit))
f.createCArray("/",
"N_k",
tables.Int64Atom(),
obj=N_k,
filters=tables.Filters(complevel=9, complib="zlib"))
f.createCArray("/",
"s_n",
tables.Int64Atom(),
obj=s_n,
filters=tables.Filters(complevel=9, complib="zlib"))
f.close()
def sample(self, N_k):
"""Draw samples from the distribution.
Parameters
----------
N_k : np.ndarray, int
number of samples per state
Returns
-------
x_kn : np.ndarray, shape=(n_states, n_samples), dtype=float
1D harmonic oscillator positions
"""
N_k = ensure_type(N_k, np.float64, 1, "N_k", self.n_states, warn_on_cast=False)
states = ["state %d" % k for k in range(self.n_states)]
x_n = []
origin_and_frame = []
for k, N in enumerate(N_k):
x0 = self.O_k[k]
sigma = (self.beta * self.K_k[k]) ** -0.5
x_n.extend(np.random.normal(loc=x0, scale=sigma, size=N))
origin_and_frame.extend([(states[k], i) for i in range(int(N))])
origin_and_frame = pd.MultiIndex.from_tuples(origin_and_frame, names=["origin", "frame"])
x_n = pd.Series(x_n, name="x", index=origin_and_frame)
u_kn = pd.DataFrame(dict([(state, x_n) for state in states]))
u_kn = 0.5 * self.K_k * (u_kn - self.O_k) ** 2.0
u_kn = u_kn.T
return x_n, u_kn, origin_and_frame
def sample(self, N_k):
"""Draw samples from the distribution.
Parameters
----------
N_k : np.ndarray, int
number of samples per state
Returns
-------
x_kn : np.ndarray, shape=(n_states, n_samples), dtype=float
1D harmonic oscillator positions
"""
N_k = ensure_type(N_k, np.float64, 1, "N_k", self.n_states, warn_on_cast=False)
states = ["state %d" % k for k in range(self.n_states)]
x_n = []
origin_and_frame = []
for k, N in enumerate(N_k):
x_n.extend(np.random.exponential(scale=self.rates[k] ** -1., size=N))
origin_and_frame.extend([(states[k], i) for i in range(int(N))])
origin_and_frame = pd.MultiIndex.from_tuples(origin_and_frame, names=["origin", "frame"])
x_n = pd.Series(x_n, name="x", index=origin_and_frame)
u_kn = np.outer(x_n, self.rates)
u_kn = pd.DataFrame(u_kn, columns=states, index=origin_and_frame).T # Note the transpose
return x_n, u_kn, origin_and_frame
def save(name, u_kn, N_k, s_n=None, least_significant_digit=None):
"""Create an HDF5 dump of an existing MBAR job for later use / testing.
Parameters
----------
name : str
Name of dataset
u_kn : np.ndarray, dtype='float', shape=(n_states, n_samples)
Reduced potential energies
N_k : np.ndarray, dtype='int', shape=(n_states)
Number of samples taken from each state
s_n : np.ndarray, optional, default=None, dtype=int, shape=(n_samples)
The state of origin of each state. If none, guess the state origins.
least_significant_digit : int, optional, default=None
If not None, perform lossy compression using tables.Filter(least_significant_digit=least_significant_digit)
Notes
-----
The output HDF5 files should be readible by the helper funtions pymbar_datasets.py
"""
import tables
(n_states, n_samples) = u_kn.shape
u_kn = ensure_type(u_kn, 'float', 2, "u_kn or Q_kn", shape=(n_states, n_samples))
N_k = ensure_type(N_k, 'int64', 1, "N_k", shape=(n_states,))
if s_n is None:
s_n = get_sn(N_k)
s_n = ensure_type(s_n, 'int64', 1, "s_n", shape=(n_samples,))
hdf_filename = os.path.join("./", "%s.h5" % name)
f = tables.File(hdf_filename, 'a')
f.createCArray("/", "u_kn", tables.Float64Atom(), obj=u_kn, filters=tables.Filters(complevel=9, complib="zlib", least_significant_digit=least_significant_digit))
f.createCArray("/", "N_k", tables.Int64Atom(), obj=N_k, filters=tables.Filters(complevel=9, complib="zlib"))
f.createCArray("/", "s_n", tables.Int64Atom(), obj=s_n, filters=tables.Filters(complevel=9, complib="zlib"))
f.close()
def __init__(self, O_k, K_k, beta=1.0):
"""Generate test case with exponential distributions.
Parameters
----------
O_k : np.ndarray, float, shape=(n_states)
Offset parameters for each state.
K_k : np.ndarray, float, shape=(n_states)
Force constants for each state.
Notes
-----
We assume potentials of the form U(x) = (k / 2) * (x - o)^2
Here, k and o are the corresponding entries of O_k and K_k.
The equilibrium distribution is given analytically by
p(x;beta,K) = sqrt[(beta K) / (2 pi)] exp[-beta K (x-x_0)**2 / 2]
The dimensionless free energy is therefore
f(beta,K) = - (1/2) * ln[ (2 pi) / (beta K) ]
"""
self.O_k = ensure_type(O_k, np.float64, 1, "O_k")
self.n_states = len(self.O_k)
self.K_k = ensure_type(K_k, np.float64, 1, "K_k", self.n_states)
self.beta = beta
def __init__(self, rates):
"""Generate test case with exponential distributions.
Parameters
----------
rates : np.ndarray, float, shape=(n_states)
Rate parameters (e.g. lambda) for each state.
Notes
-----
We assume potentials of the form U(x) = lambda x.
"""
rates = ensure_type(rates, np.float64, 1, "rates")
self.n_states = len(rates)
self.rates = rates
