Skip to content

Commit e4c7c0e

Browse files
xiki-tempulaxiki-tempula
committed
update
1 parent 769b156 commit e4c7c0e

6 files changed

Lines changed: 252 additions & 8 deletions

File tree

‎.gitignore‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -159,4 +159,5 @@ cython_debug/
159159
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
160160
#.idea/
161161

162-
.DS_Store
162+
.DS_Store
163+
.idea/

‎src/bayesmbar/__init__.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,5 +2,6 @@
22
from .cbayesmbar import CBayesMBAR
33
from .bayesbar import BayesBAR
44
from .fastmbar import FastMBAR
5+
from .offsetmbar import OffsetMBAR
56

6-
__all__ = ('FastMBAR', 'BayesMBAR', 'BayesBAR', 'CBayesMBAR', '__init__')
7+
__all__ = ('FastMBAR', 'BayesMBAR', 'BayesBAR', 'CBayesMBAR', 'OffsetMBAR', '__init__')

‎src/bayesmbar/bayesmbar.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -119,7 +119,8 @@ def logdensity(dF):
119119

120120
## compute the mean, covariance, and precision of dF based on the samples from the likelihood
121121
self._dF_mean_ll = jnp.mean(self._dF_samples_ll, axis=0)
122-
self._dF_cov_ll = jnp.cov(self._dF_samples_ll.T)
122+
# Ensure covariance is 2D for Cholesky decomposition (handle scalar case)
123+
self._dF_cov_ll = jnp.atleast_2d(jnp.cov(self._dF_samples_ll.T))
123124

124125
L = jnp.linalg.cholesky(self._dF_cov_ll)
125126
L_inv = jax.scipy.linalg.solve_triangular(L, jnp.eye(L.shape[0]), lower=True)

‎src/bayesmbar/offsetmbar.py‎

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
import numpy as np
2+
3+
from bayesmbar import CBayesMBAR
4+
from numpy import ndarray
5+
6+
7+
class OffsetMBAR:
8+
def __init__(
9+
self,
10+
energies: list[np.ndarray],
11+
nums_conf: list[np.ndarray],
12+
offsets: list[float],
13+
sample_size: int = 1000,
14+
warmup_steps: int = 500,
15+
method: str = "Newton",
16+
random_seed: int = None,
17+
verbose: bool = True,
18+
) -> None:
19+
"""
20+
Offset-constrained Coupled BayesMBAR.
21+
22+
This variant of Coupled BayesMBAR enforces a simple constraint: for each
23+
coupled system i, the free-energy difference computed from its energies
24+
plus a provided scalar offset must be identical across all systems. In
25+
other words, the solution to the MBAR equation satisfies
26+
mbar(energies[i]) + offsets[i] = constant for every system i.
27+
28+
Parameters
29+
----------
30+
energies : List[numpy.ndarray]
31+
Per-system arrays of reduced potentials (in units of kT) used by
32+
MBAR. One array per coupled system.
33+
nums_conf : List[numpy.ndarray]
34+
Per-system arrays of counts (number of configurations) corresponding
35+
to the provided energies.
36+
offsets : List[float]
37+
One scalar offset per system. The constraint is that
38+
mbar(energies[i]) + offsets[i] is the same across all systems.
39+
sample_size : int, optional
40+
Number of samples to draw from the likelihood. Default: 1000.
41+
warmup_steps : int, optional
42+
Number of warmup steps for the HMC sampler. Default: 500.
43+
method : str, optional
44+
Optimization method to find the likelihood mode. Either "Newton" or
45+
"L-BFGS-B". Default: "Newton".
46+
random_seed : int, optional
47+
Random seed. If None, a seed is generated from the current time.
48+
Default: None.
49+
verbose : bool, optional
50+
If True, print sampling progress. Default: True.
51+
"""
52+
new_energies = []
53+
new_nums_conf = []
54+
first_state = []
55+
last_state = []
56+
connecting_states = []
57+
index = 0
58+
for energy, n_conf, offset in zip(energies, nums_conf, offsets, strict=True):
59+
# Add the original reduced potential
60+
states, n_samples = energy.shape
61+
new_energies.append(energy)
62+
new_nums_conf.append(n_conf)
63+
first_state.append((index, 0))
64+
# Generate the reduced potential for offset
65+
index += 1
66+
# Make it divisible
67+
slice = energy[0, :n_samples]
68+
new_energy = np.linspace(0, offset, states).reshape(
69+
(states, 1)
70+
) + slice.reshape((1, n_samples))
71+
new_energies.append(new_energy)
72+
new_nums_conf.append(n_conf)
73+
last_state.append((index, states - 1))
74+
connecting_states.append([(index - 1, states - 1), (index, 0)])
75+
index += 1
76+
77+
self.cbmbar = CBayesMBAR(
78+
new_energies,
79+
new_nums_conf,
80+
identical_states=[first_state, last_state, *connecting_states],
81+
method=method,
82+
sample_size=sample_size,
83+
warmup_steps=warmup_steps,
84+
random_seed=random_seed,
85+
verbose=verbose,
86+
)
87+
88+
@property
89+
def F_mode(self) -> list[ndarray]:
90+
F_mode_list = []
91+
for real_f, offset_f in zip(self.cbmbar.F_mode[::2], self.cbmbar.F_mode[1::2]):
92+
F_mode = np.append(real_f, real_f[-1] + offset_f[-1] - offset_f[0])
93+
F_mode_list.append(F_mode)
94+
return F_mode_list
95+
96+
@property
97+
def F_mean(self) -> list[ndarray]:
98+
F_mean_list = []
99+
for real_f, offset_f in zip(self.cbmbar.F_mean[::2], self.cbmbar.F_mean[1::2]):
100+
F_mean = np.append(real_f, real_f[-1] + offset_f[-1] - offset_f[0])
101+
F_mean_list.append(F_mean)
102+
return F_mean_list

‎src/test/test_bayesmbar.py‎

Lines changed: 40 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
1+
import numpy as np
12
import pytest
23
from pytest import approx
34
from bayesmbar import BayesMBAR
45

6+
57
@pytest.mark.parametrize("method", ["Newton", "L-BFGS-B"])
68
def test_BayesMBAR(setup_mbar_data, method):
79
energy, num_conf, F_ref, energy_p, F_ref_p = setup_mbar_data
@@ -11,12 +13,45 @@ def test_BayesMBAR(setup_mbar_data, method):
1113
verbose=True,
1214
method=method,
1315
)
14-
assert mbar.F_mode == approx(F_ref, abs = 1e-1)
15-
assert mbar.F_mean == approx(F_ref, abs = 1e-1)
16+
assert mbar.F_mode == approx(F_ref, abs=1e-1)
17+
assert mbar.F_mean == approx(F_ref, abs=1e-1)
18+
19+
20+
def test_two_states():
21+
M = 2 ## number of states
22+
mu = np.linspace(0, 1, M) ## equilibrium positions
23+
k = np.random.uniform(10, 30, M) ## force constants
24+
sigma = np.sqrt(1.0 / k)
25+
F_reference = -np.log(sigma)
26+
F_reference -= F_reference[0]
27+
n = 100
28+
x = [np.random.normal(mu[i], sigma[i], (n,)) for i in range(M)]
29+
x = np.concatenate(x)
30+
u = 0.5 * k.reshape((-1, 1)) * (x - mu.reshape((-1, 1))) ** 2
31+
num_conf = np.array([n for i in range(M)])
32+
mbar = BayesMBAR(
33+
u,
34+
num_conf,
35+
prior="uniform",
36+
mean=None,
37+
state_cv=None,
38+
kernel=None,
39+
sample_size=1000,
40+
warmup_steps=100,
41+
optimize_steps=0,
42+
random_seed=0,
43+
verbose=False,
44+
)
45+
F_reference = F_reference[-1] - F_reference[0]
46+
F_mean = mbar.F_mean
47+
F_mode = mbar.F_mode
48+
assert (F_mean[-1] - F_mean[0]) == approx(F_reference, abs=1)
49+
assert (F_mode[-1] - F_mode[0]) == approx(F_reference, abs=1)
50+
51+
# results = fastmbar.calculate_free_energies_of_perturbed_states(energy_p)
52+
# results['F'] = results['F'] - results['F'].mean()
53+
# assert results['F'] == approx(F_ref_p, abs = 1e-1)
1654

17-
#results = fastmbar.calculate_free_energies_of_perturbed_states(energy_p)
18-
#results['F'] = results['F'] - results['F'].mean()
19-
#assert results['F'] == approx(F_ref_p, abs = 1e-1)
2055

2156
# import pytest
2257
# from pytest import approx

‎src/test/test_offsetmbar.py‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
import numpy as np
2+
import pytest
3+
from bayesmbar import CBayesMBAR, OffsetMBAR
4+
import numpy.testing as npt
5+
6+
n_samples = 100
7+
8+
9+
@pytest.fixture()
10+
def no_offset():
11+
M = 5 ## number of states
12+
repeats = 3
13+
mu = np.linspace(0, 1, M) ## equilibrium positions
14+
k = np.random.uniform(10, 30, M) ## force constants
15+
sigma = np.sqrt(1.0 / k)
16+
F_reference = -np.log(sigma)
17+
F_reference -= F_reference[0]
18+
n = n_samples
19+
energies = []
20+
for _ in range(repeats):
21+
x = [np.random.normal(mu[i], sigma[i], (n,)) for i in range(M)]
22+
x = np.concatenate(x)
23+
u = 0.5 * k.reshape((-1, 1)) * (x - mu.reshape((-1, 1))) ** 2
24+
energies.append(u)
25+
num_conf = [np.array([n for i in range(M)])] * repeats
26+
return F_reference, energies, num_conf, repeats
27+
28+
29+
def test_cmbar(no_offset):
30+
F_reference, energies, num_conf, repeats = no_offset
31+
cmbar = CBayesMBAR(
32+
energies,
33+
num_conf,
34+
identical_states=[
35+
[(i, 0) for i in range(repeats)],
36+
[(i, 4) for i in range(repeats)],
37+
],
38+
sample_size=1000,
39+
warmup_steps=100,
40+
random_seed=0,
41+
verbose=False,
42+
)
43+
f_mean = cmbar.F_mean
44+
45+
assert f_mean[0][-1] == pytest.approx(F_reference[-1], abs=0.2)
46+
f = [d_f[-1] for d_f in f_mean]
47+
npt.assert_allclose(f, f[0], atol=1e-5)
48+
49+
50+
def test_0offset(no_offset):
51+
F_reference, energies, num_conf, repeats = no_offset
52+
cmbar = OffsetMBAR(
53+
energies,
54+
num_conf,
55+
offsets=[0] * repeats,
56+
sample_size=1000,
57+
warmup_steps=100,
58+
random_seed=0,
59+
verbose=False,
60+
)
61+
62+
f_mode = cmbar.F_mode
63+
f = [d_f[-1] for d_f in f_mode]
64+
npt.assert_allclose(f, f[0], atol=0.1)
65+
for edge in f_mode:
66+
assert edge[-1] - edge[-2] == pytest.approx(0, abs=0.1)
67+
assert f_mode[0][-1] == pytest.approx(F_reference[-1], abs=0.2)
68+
69+
70+
def test_offset():
71+
M = 5 ## number of states
72+
repeats = 3
73+
mu = np.linspace(0, 1, M) ## equilibrium positions
74+
n = n_samples
75+
F_reference_list = []
76+
energies = []
77+
for _ in range(repeats):
78+
k = np.random.uniform(10, 30, M) ## force constants
79+
sigma = np.sqrt(1.0 / k)
80+
F_reference = -np.log(sigma)
81+
F_reference -= F_reference[0]
82+
F_reference_list.append(F_reference)
83+
84+
x = [np.random.normal(mu[i], sigma[i], (n,)) for i in range(M)]
85+
x = np.concatenate(x)
86+
u = 0.5 * k.reshape((-1, 1)) * (x - mu.reshape((-1, 1))) ** 2
87+
energies.append(u)
88+
offsets = [-1 * F_reference[-1] for F_reference in F_reference_list]
89+
num_conf = [np.array([n for i in range(M)])] * repeats
90+
cmbar = OffsetMBAR(
91+
energies,
92+
num_conf,
93+
offsets=offsets,
94+
sample_size=1000,
95+
warmup_steps=100,
96+
random_seed=0,
97+
verbose=False,
98+
)
99+
100+
f_mode = cmbar.F_mode
101+
for edge, offset in zip(f_mode, offsets):
102+
assert edge[0] == pytest.approx(0, abs=0.1)
103+
assert edge[-1] == pytest.approx(0, abs=0.1)
104+
assert edge[-1] - edge[-2] == offset

0 commit comments

Comments
 (0)