Solving a Real Problem Using Our Model: The Discrete Flashing Ratchet (DFR)

The following is adapted from Berghaus et al. (2024) and shows a step by step example for solving problems using our trained model from Hugging Face. This notebook is a detailed example of solving a physics problem using our model using the following approach:

Understanding the problem
Loading and understanding the data
Using our model to infer the transition rates
Using the inferred transition rates to solve the original problem

In statistical physics, the ratchet effect refers to the rectification of thermal fluctuations into directed motion to produce work, and goes all the way back to Feynman Feynman et al. (1965).

Here we consider a simple example thereof, in which a Brownian particle, immersed in a thermal bath at unit temperature, moves on a one-dimensional lattice. The particle is subject to a linear, periodic and asymmetric potential of maximum height $2V$ that is switched on and off at a constant rate $r$ . The potential has three possible values when is switched on, which correspond to three of the states of the system. The particle jumps among them with rate $f_{ij}^{\tiny{\text{ on}}}$ .

When the potential is switched off, the particle jumps freely with rate $f_{ij}^{\tiny{\text{ off}}}$ . We can therefore think of the system as a six-state system, as illustrated here:

Similar to Roldán & Parrondo (2010), we now define the transition rates as

f_{ij}^{\tiny{\text{on}}} = \exp \left( -\frac{V}{2}(j-i)\right), \, \, \, \text{for} \, \, i, j \in (0, 1, 2); \quad f_{ij}^{\tiny{\text{ off}}} = b, \, \, \, \text{for} \, \, i, j \in (3, 4, 5).

(1)

Given these specifics, we consider the parameter set $(V, r, B) = (1, 1, 1)$
together with the dataset simulated by Seifner & Sánchez (2023),

So given our data, we want to recover the theoretical transition rates, as a q-matrix.^[1]

import warnings
warnings.filterwarnings("ignore")
import numpy as np
np.set_printoptions(linewidth=100)
from transformers.utils import logging
logging.disable_progress_bar()
from datasets.utils.logging import disable_progress_bar 
disable_progress_bar()

rates_on = np.array([[np.exp(-0.5 * (j - i)) - 1 * (i == j) for j in range(3)] for i in range(3)])
rates_off = np.ones(shape=(3, 3))
np.fill_diagonal(rates_off, 0)
id3 = np.eye(3)
q_matrix = np.zeros(shape=(6, 6))
q_matrix[:3, :3] = rates_on
q_matrix[3:, 3:] = rates_off
q_matrix[3:, :3] = id3
q_matrix[:3, 3:] = id3
diagonal = -np.sum(q_matrix, axis=1)
np.fill_diagonal(q_matrix, diagonal)
q_matrix

array([[-1.9744101 ,  0.60653066,  0.36787944,  1.        ,  0.        ,  0.        ],
       [ 1.64872127, -3.25525193,  0.60653066,  0.        ,  1.        ,  0.        ],
       [ 2.71828183,  1.64872127, -5.3670031 ,  0.        ,  0.        ,  1.        ],
       [ 1.        ,  0.        ,  0.        , -3.        ,  1.        ,  1.        ],
       [ 0.        ,  1.        ,  0.        ,  1.        , -3.        ,  1.        ],
       [ 0.        ,  0.        ,  1.        ,  1.        ,  1.        , -3.        ]])

Loading and Exploring the Data¶

We start by loading the simulated data from Hugging Face:

# Loading the Discrete Flashing Ratchet (DFR) dataset from Huggingface
import torch
from datasets import load_dataset

data = load_dataset("FIM4Science/mjp", download_mode="force_redownload", name="default")
data.set_format("torch")

Repo card metadata block was not found. Setting CardData to empty.

The model will later use the ‘observation_grid’, ‘observation_values’, ‘seq_lengths’ features to estimate the transition rates.

Here the observation grid is constant over all paths, consisting of 100 evenly spaced points on $[0,1]$ . The sequence lengths of the dataset is therefore just 100 for all paths.

data["train"]["observation_grid"][0][0, 0, :, 0]
print(data["train"]["observation_values"][0][0, 0, :, 0])
print(data["train"]["observation_values"][0][0, 300, :, 0])

tensor([2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 4, 4, 4, 4, 4, 2, 2, 2,
        2, 2, 2, 2, 2, 2, 2, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 3, 3, 3, 3, 3, 3,
        3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4,
        4, 4, 4, 4])
tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 3, 3, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
        2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
        2, 2, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5,
        5, 5, 5, 5])

The observation values contain the state of the processes for all paths and all time points.

We can visualize the first 3 paths to get a feeling for the processes:

import matplotlib.pyplot as plt
import matplotlib.colors as mcolors

cmap = mcolors.LinearSegmentedColormap.from_list("custom_cmap", ["#D55E00", "#56B4E9"], N=3)
colors = [mcolors.to_hex(cmap(i)) for i in range(3)]

ts = data["train"]["observation_grid"][0][0, 0, :, 0]
three_paths = data["train"]["observation_values"][0][0, :3, :, 0]
for i in range(3):
    plt.plot(ts, three_paths[i], "x--", c=colors[i])
plt.show()

Inferring Transition Rates¶

To infer the $q$ -matrix we first load our trained model:

from fim.trainers.utils import get_accel_type
from fim.models.mjp import FIMMJP

fimmjp = FIMMJP.from_pretrained("FIM4Science/fim-mjp", trust_remote_code=True)

device = get_accel_type()
fimmjp = fimmjp.to(device)
fimmjp.eval()

FIMMJP(
  (gaussian_nll): GaussianNLLLoss()
  (init_cross_entropy): CrossEntropyLoss()
  (pos_encodings): SineTimeEncoding(
    (linear_embedding): Linear(in_features=1, out_features=1, bias=True)
    (periodic_embedding): Sequential(
      (0): Linear(in_features=1, out_features=249, bias=True)
      (1): SinActivation()
    )
  )
  (ts_encoder): TransformerEncoder(
    (layers): ModuleList(
      (0-3): 4 x TransformerEncoderLayer(
        (self_attn): MultiheadAttention(
          (out_proj): NonDynamicallyQuantizableLinear(in_features=256, out_features=256, bias=True)
        )
        (linear1): Linear(in_features=256, out_features=1024, bias=True)
        (dropout): Dropout(p=0.1, inplace=False)
        (linear2): Linear(in_features=1024, out_features=256, bias=True)
        (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True, bias=True)
        (norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True, bias=True)
        (dropout1): Dropout(p=0.1, inplace=False)
        (dropout2): Dropout(p=0.1, inplace=False)
      )
    )
  )
  (path_attention): MultiHeadLearnableQueryAttention(
    (W_k): Linear(in_features=256, out_features=256, bias=False)
    (W_v): Linear(in_features=256, out_features=256, bias=False)
    (W_o): Linear(in_features=256, out_features=256, bias=False)
  )
  (intensity_matrix_decoder): MLP(
    (layers): Sequential(
      (linear_0): Linear(in_features=257, out_features=128, bias=True)
      (activation_0): SELU()
      (dropout_0): Dropout(p=0.1, inplace=False)
      (linear_1): Linear(in_features=128, out_features=128, bias=True)
      (activation_1): SELU()
      (dropout_1): Dropout(p=0.1, inplace=False)
      (output): Linear(in_features=128, out_features=60, bias=True)
    )
  )
  (initial_distribution_decoder): MLP(
    (layers): Sequential(
      (linear_0): Linear(in_features=257, out_features=128, bias=True)
      (activation_0): SELU()
      (dropout_0): Dropout(p=0.1, inplace=False)
      (linear_1): Linear(in_features=128, out_features=128, bias=True)
      (activation_1): SELU()
      (dropout_1): Dropout(p=0.1, inplace=False)
      (output): Linear(in_features=128, out_features=6, bias=True)
    )
  )
)

As noted in Berghaus et al. (2024)^[2] it suffices to look at a small context window of 300 paths with 50 observation values each. We therefore infer the transition rates batchwise to demonstrate how little data might be needed.

# copy data to device
batch = {
    k: v.to(device)[0]
    for k, v in data["train"][:1].items()
    if k not in ["intensity_matrices", "adjacency_matrices", "initial_distributions"]
}  # data without any information, we seek to find

n_paths = 300
total_n_paths = batch["observation_grid"][0].shape[0]
statistics = 50

outputs = []

with torch.no_grad():
    for _ in range(statistics):
        paths_idx = torch.randperm(total_n_paths)[:n_paths]
        mini_batch = batch.copy()
        mini_batch["observation_grid"] = batch["observation_grid"][:, paths_idx]
        mini_batch["observation_values"] = batch["observation_values"][:, paths_idx]
        mini_batch["seq_lengths"] = batch["seq_lengths"][:, paths_idx]
        
        outputs.append(
            fimmjp(mini_batch, n_states=6)["intensity_matrices"]
        )  # We are currently not interested in variances or initial conditions

This results in the following inferred transition rates:

mean_rates = torch.mean(torch.stack(outputs, dim=0), dim=0)[0].to("cpu").numpy()
print(mean_rates)
plt.matshow(mean_rates)
plt.colorbar()
plt.show()

[[-2.2008178   0.59146845  0.3066814   1.1464729   0.08717554  0.0690197 ]
 [ 1.8060498  -3.621488    0.56928617  0.1142701   1.02205     0.10983212]
 [ 2.4457183   1.9562448  -5.8043838   0.2503286   0.21888165  0.93321043]
 [ 0.9128032   0.10649365  0.08033262 -3.2453835   0.99051064  1.1552435 ]
 [ 0.08385407  0.9179251   0.12883677  1.0382289  -3.2431674   1.0743225 ]
 [ 0.13743013  0.16307525  0.9567804   1.0563439   0.9520308  -3.2656603 ]]

and the following variances:

variances = torch.var(torch.stack(outputs, dim=0), dim=0)[0].to("cpu").numpy()
print(variances)
plt.matshow(variances)
plt.colorbar()
plt.show()

[[1.9652484e-02 5.5053513e-03 2.8013338e-03 8.0227237e-03 1.7621358e-05 1.4784444e-05]
 [1.7133208e-02 6.6398464e-02 2.4012363e-02 5.5016371e-05 1.4633543e-02 9.3533068e-05]
 [4.0423211e-02 4.5272861e-02 1.0089903e-01 7.4126315e-04 3.2079371e-04 3.3523783e-02]
 [5.8041019e-03 1.8132603e-05 5.4408003e-05 2.7347988e-02 9.6584782e-03 7.8097358e-03]
 [2.6498530e-05 1.2430202e-02 9.0584705e-05 6.5338137e-03 2.2267068e-02 1.1575093e-02]
 [2.2110729e-04 1.8194533e-04 2.2850266e-02 1.4584687e-02 2.1373078e-02 6.2594831e-02]]

Using the previously computed theoretical transition rates we can look at the pointwise error:

error = mean_rates - q_matrix
print(error)
plt.matshow(error)
plt.colorbar()
plt.show()

[[-0.22640772 -0.01506221 -0.06119805  0.14647293  0.08717554  0.0690197 ]
 [ 0.15732855 -0.36623616 -0.03724449  0.1142701   0.02205002  0.10983212]
 [-0.27256354  0.30752356 -0.43738066  0.2503286   0.21888165 -0.06678957]
 [-0.08719683  0.10649365  0.08033262 -0.2453835  -0.00948936  0.15524352]
 [ 0.08385407 -0.08207488  0.12883677  0.03822887 -0.2431674   0.07432246]
 [ 0.13743013  0.16307525 -0.04321963  0.05634391 -0.04796922 -0.26566029]]

Inferring the Parameters¶

We will now use the transition rates to recover the parameters used to generate the dataset. Recall the formula we previously used to calculate the theoretical transition rates:

f_{ij}^{\tiny{\text{on}}} = \exp \left( -\frac{V}{2}(j-i)\right), \, \, \, \text{for} \, \, i, j \in (0, 1, 2); \quad f_{ij}^{\tiny{\text{ off}}} = b, \, \, \, \text{for} \, \, i, j \in (3, 4, 5).

(2)

The $q$ -matrix (transition rate matrix) is therefore given by

q=\begin{pmatrix} f_{ij}^{\tiny{\text{on}}} & r\cdot I_3\\ r\cdot I_3 & f_{ij}^{\tiny{\text{ off}}} \end{pmatrix}

(3)

since the transition rates between $(i,j,\text{on})$ and $(i,j,\text{off})$ (and vice versa) are given by the parameter $r$ .

This system is evidently overdetermined. For simplicity we assume normally distributed independent errors, which results in the following MLEs for each of the parameters:

We can easily recover $b$ using the diagonal:

\hat{b}=\frac{1}{6}\sum_{i=3}^5\sum_{j=3}^51_{i\neq j}f_{ij}^{\tiny{\text{off}}}

(4)

lower_matrix = mean_rates[3:, 3:]
np.fill_diagonal(lower_matrix, 0)

B_hat = lower_matrix.sum() / 6
B_hat

np.float32(1.0444467)

Similarly $r$ describes the transition rates from a state $(i,j,\text{on})$ to $(i,j,\text{off})$ and vice versa:

\hat{r}=\frac{1}{6}\sum_{i=0}^2 q_{i,i+3}+q_{i+3,i}

(5)

r_hat = sum([mean_rates[i, i + 3] + mean_rates[i + 3, i] for i in range(3)]) / 6
r_hat

np.float32(0.9815404)

And $V$ can be recovered using $(q_{i,j})_{i,j=1}^3$ :

\hat{V}=\frac{1}{6}\left(-2\cdot(\ln q_{0,1}+\ln q_{1,2})+2\cdot(\ln q_{1,0}+\ln q_{2,1})- \ln q_{0,2}+ \ln q_{2,0}\right)

(6)

V_hat = (
    -2 * (np.log(mean_rates[0, 1]) + np.log(mean_rates[1, 2]))
    + 2 * (np.log(mean_rates[1, 0]) + np.log(mean_rates[2, 1]))
    - np.log(mean_rates[0, 2])
    + np.log(mean_rates[2, 0])
) / 6
V_hat

np.float32(1.12961)

Which results in the following inferred parameters:

(V_hat, r_hat, B_hat)

(np.float32(1.12961), np.float32(0.9815404), np.float32(1.0444467))

Footnotes¶

That is with the diagonal entries normalized to the negative sum of the transition rates
↩
Section 4.1 and E.4
↩

References¶

Berghaus, D., Cvejoski, K., Seifner, P., Ojeda, C., & Sánchez, R. J. (2024). Foundation Inference Models for Markov Jump Processes. The Thirty-Eighth Annual Conference on Neural Information Processing Systems. https://openreview.net/forum?id=f4v7cmm5sC
Feynman, R. P., Leighton, R. B., Sands, M., & Hafner, E. M. (1965). The feynman lectures on physics; vol. i. American Journal of Physics, 33(9), 750–752.
Roldán, É., & Parrondo, J. M. R. (2010). Estimating dissipation from single stationary trajectories. Physical Review Letters, 105 15, 150607.
Seifner, P., & Sánchez, R. J. (2023). Neural Markov Jump Processes. International Conference on Machine Learning, 30523–30552. https://proceedings.mlr.press/v202/seifner23a/seifner23a.pdf