Update README.md

README.md

@@ -1,201 +1,260 @@
 ---
-library_name: transformers
+tags:
+- sequence-classification
+- time-series
+- stochastic-processes
+- markov-jump-processes
 license: cc-by-4.0
 datasets:
-- cvejoski/MJP
+- custom
+metrics:
+- rmse
+- Hellinger distance
 ---
 
-# Model Card for Model ID
+# Foundation Inference Model (FIM) for Markov Jump Processes  Model Card
 
-<!-- Provide a quick summary of what the model is/does. -->
+## Model Description
 
+The Foundation Inference Model (`FIM`) is a neural recognition model designed for zero-shot inference of Markov Jump Processes (MJPs) in bounded state spaces. FIM processes noisy and sparse observations to estimate the transition rate matrix and initial condition of MJPs, without requiring fine-tuning on the target dataset.
 
+FIM combines supervised learning on a synthetic dataset of MJPs with attention mechanisms, enabling robust inference for empirical processes with varying dimensionalities. It is the first generative zero-shot model for MJPs, offering broad applicability across domains such as molecular dynamics, ion channel dynamics, and discrete flashing ratchet systems.
 
-## Model Details
+[![GitHub](https://img.shields.io/badge/GitHub-Repository-blue?logo=github)](https://github.com/cvejoski/OpenFIM)
+[![arXiv](https://img.shields.io/badge/arXiv-2406.06419-B31B1B.svg)](https://arxiv.org/abs/2406.06419)
 
-### Model Description
+## Intended Use
 
-<!-- Provide a longer summary of what this model is. -->
+- Applications:
+   - Inferring dynamics of physical, chemical, and biological systems.
+   - Estimating transition rates and initial conditions from noisy observations.
+   - Zero-shot simulation and analysis of MJPs for:
+      - Molecular dynamics simulations (e.g., alanine dipeptide conformations).
+      - Protein folding models.
+      - Ion channel dynamics.
+      - Brownian motor systems.
+- Users: Researchers in statistical physics, molecular biology, and stochastic processes.
 
-This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
+- Limitations:
+   - The model performs well only for processes with dynamics similar to its synthetic training distribution.
+   - Poor estimates are likely for datasets with distributions significantly deviating from the synthetic priors (e.g., systems with power-law distributed rates).
 
-- **Developed by:** [More Information Needed]
-- **Funded by [optional]:** [More Information Needed]
-- **Shared by [optional]:** [More Information Needed]
-- **Model type:** [More Information Needed]
-- **Language(s) (NLP):** [More Information Needed]
-- **License:** [More Information Needed]
-- **Finetuned from model [optional]:** [More Information Needed]
+### Installation
 
-### Model Sources [optional]
+To install the Foundation Inference Model (FIM) from the fim library, follow these steps:
 
-<!-- Provide the basic links for the model. -->
+1. **Clone the repository:**
 
-- **Repository:** [More Information Needed]
-- **Paper [optional]:** [More Information Needed]
-- **Demo [optional]:** [More Information Needed]
+   ```bash
+   git clone https://github.com/cvejoski/OpenFIM.git
+   cd OpenFIM
+   ```
 
-## Uses
+2. **Install the required dependencies:**
 
-<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+   ```bash
+   pip install -r requirements.txt
+   ```
 
-### Direct Use
+3. **Install the fim library:**
 
-<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+   ```bash
+   pip install .
+   ```
 
-[More Information Needed]
+After completing these steps, you should have the `fim` library installed and ready to use.
 
-### Downstream Use [optional]
+### Description of the Input and Output Data of the Forward Step of the FIMMJP Model
 
-<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+#### Input Data
 
-[More Information Needed]
+The input data to the forward step of the `FIMMJP` model is a dictionary containing several key-value pairs. Each key corresponds to a specific type of input data required by the model. Below is a detailed description of each key and its corresponding value:
 
-### Out-of-Scope Use
+1. **`observation_grid`**:
+   - **Type**: `torch.Tensor`
+   - **Shape**: `[B, P, L, 1]`
+   - **Description**: This tensor represents the observation grid. `B` is the batch size, `P` is the number of paths, `L` is the length of each path, and `1` indicates a single time dimension.
 
-<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+2. **`observation_values`**:
+   - **Type**: torch.Tensor
+   - **Shape**: `[B, P, L, D]`
+   - **Description**: This tensor contains the observation values. `D` is the dimensionality of the observations.
 
-[More Information Needed]
+3. **`seq_lengths`**:
+   - **Type**: torch.Tensor
+   - **Shape**: `[B, P]`
+   - **Description**: This tensor represents the sequence lengths for each path in the batch.
+4. **`initial_distributions`**:
+     - **Type**: torch.Tensor
+     - **Shape**: `[B, N]`
+     - **Description**: This tensor represents the initial distributions.
 
-## Bias, Risks, and Limitations
+4. **Optional Keys**:
+   - **`time_normalization_factors`**:
+     - **Type**: torch.Tensor
+     - **Shape**: `[B, 1]`
+     - **Description**: This tensor represents the time normalization factors.
+   - **`intensity_matrices`**:
+     - **Type**: torch.Tensor
+     - **Shape**: `[B, N, N]`
+     - **Description**: This tensor represents the intensity matrices.
 
-<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+   - **`adjacency_matrices`**:
+     - **Type**: torch.Tensor
+     - **Shape**: `[B, N, N]`
+     - **Description**: This tensor represents the adjacency matrices.
 
-[More Information Needed]
+#### Output Data
 
-### Recommendations
+The output data from the forward step of the `FIMMJP` model is a dictionary containing the following key-value pairs:
 
-<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+1. **`intensity_matrices`**:
+   - **Type**: torch.Tensor
+   - **Shape**: `[B, N, N]`
+   - **Description**: This tensor represents the predicted intensity matrix for each sample in the batch. `N` is the number of states in the process.
 
-Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+2. **`intensity_matrices_variance`**:
+   - **Type**: torch.Tensor
+   - **Shape**: `[B, N, N]`
+   - **Description**: This tensor represents the log variance of the predicted intensity matrix for each sample in the batch.
 
-## How to Get Started with the Model
+3. **`initial_condition`**:
+   - **Type**: torch.Tensor
+   - **Shape**: `[B, N]`
+   - **Description**: This tensor represents the predicted initial distribution of states for each sample in the batch.
 
-Use the code below to get started with the model.
+4. **`losses`** (optional):
+   - **Type**: dict
+   - **Description**: This dictionary contains the calculated losses if the required keys (`intensity_matrices` and `initial_distributions`) are present in the input data. The keys in this dictionary include:
+     - **loss**: The total loss.
+     - **loss_gauss**: The Gaussian negative log-likelihood loss.
+     - **loss_initial**: The cross-entropy loss for the initial distribution.
+     - **loss_missing_link**: The loss for missing links in the intensity matrix.
+     - **rmse_loss**: The root mean square error loss.
+     - **`beta_gauss_nll`**: The weight for the Gaussian negative log-likelihood loss.
+     - **`beta_init_cross_entropy`**: The weight for the cross-entropy loss.
+     - **`beta_missing_link`**: The weight for the missing link loss.
+     - **`number_of_paths`**: The number of paths in the batch.
 
-[More Information Needed]
+### Example Usage
 
-## Training Details
+Here is an example of how to use the `FIMMJP` model for inference:
 
-### Training Data
+```python
+import torch
+from transformers import AutoModel
 
-<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-[More Information Needed]
+# Loading the model
+model = AutoModel.from_pretrained("cvejoski/FIMMJP", trust_remote_code=True)
+model = model.to(device)
+model.eval()
 
-### Training Procedure
+# Loading the Discrete Flashing Ratchet (DFR) dataset from Huggingface
+data = load_dataset("cvejoski/mjp", download_mode="force_redownload", trust_remote_code=True, name="DFR_V=1")
+data.set_format("torch")
 
-<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+# Create batch
+inputs = {k: v.to(device) for k, v in data["train"][:1].items()}
 
-#### Preprocessing [optional]
+# Perform inference
+outputs = model(inputs, n_states=6)
 
-[More Information Needed]
+# Process the output as needed
+intensity_matrix = outputs["intensity_matrices"]
+initial_distribution = outputs["initial_condition"]
 
+print(intensity_matrix)
+print(initial_distribution)
+```
 
-#### Training Hyperparameters
+In this example, the input data is prepared and passed to the model's forward step. The model returns the predicted intensity matrix and initial distribution, which can then be processed as needed.
 
-- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+## Model Training
 
-#### Speeds, Sizes, Times [optional]
+### Training Dataset:
 
-<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+- Synthetic MJPs covering state spaces ranging from 2 to 6 states, with up to 300 paths per process.
+- Training spans 45,000 MJPs sampled using the Gillespie algorithm, with various grid and noise configurations.
+- Noise: Includes mislabeled states (1% to 10% noise).
+- Observations: Regular and irregular grids with up to 100 time points.
 
-[More Information Needed]
 
-## Evaluation
-
-<!-- This section describes the evaluation protocols and provides the results. -->
-
-### Testing Data, Factors & Metrics
-
-#### Testing Data
-
-<!-- This should link to a Dataset Card if possible. -->
-
-[More Information Needed]
-
-#### Factors
-
-<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
-
-[More Information Needed]
-
-#### Metrics
-
-<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+### Architecture:
 
-[More Information Needed]
+- `Input`: K time series of noisy observations and their associated grids.
+- `Encoder`: LSTM or Transformer for time-series embedding.
+- `Attention`: Self-attention mechanism aggregates embeddings.
+- `Output`: Transition rate matrix, variance matrix, and initial distribution.
 
-### Results
+### Loss Function:
 
-[More Information Needed]
+- Supervised likelihood maximization, with regularization for missing links in the intensity matrix.
 
-#### Summary
-
-
-
-## Model Examination [optional]
-
-<!-- Relevant interpretability work for the model goes here -->
-
-[More Information Needed]
-
-## Environmental Impact
-
-<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
-
-Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+## Evaluation
 
-- **Hardware Type:** [More Information Needed]
-- **Hours used:** [More Information Needed]
-- **Cloud Provider:** [More Information Needed]
-- **Compute Region:** [More Information Needed]
-- **Carbon Emitted:** [More Information Needed]
+The Foundation Inference Model (FIM) was evaluated on a diverse set of datasets to demonstrate its zero-shot inference capabilities for Markov Jump Processes (MJPs). The evaluation spans datasets representing different domains, such as statistical physics, molecular dynamics, and experimental biological data, testing FIM's ability to infer transition rate matrices, initial distributions, and compute physical properties like stationary distributions and relaxation times.
 
-## Technical Specifications [optional]
+### Datasets
 
-### Model Architecture and Objective
+The following datasets were used to evaluate FIM:
 
-[More Information Needed]
+1. Discrete Flashing Ratchet (DFR):
 
-### Compute Infrastructure
+   - A 6-state stochastic model of a Brownian motor under a periodic potential.
+   - Dataset: 5,000 paths recorded on an irregular grid of 50 time points.
+   - Metrics: Transition rates, stationary distributions, and entropy production.
 
-[More Information Needed]
+2. Switching Ion Channel (IonCh):
 
-#### Hardware
+   - A 3-state model of ion flow across viral potassium channels.
+   - Dataset: Experimental recordings of 5,000 paths sampled at 5kHz over one second.
+   - Metrics: Mean first-passage times, stationary distributions.
 
-[More Information Needed]
+3. Alanine Dipeptide (ADP):
 
-#### Software
+   - A 6-state molecular dynamics model describing dihedral angles of alanine dipeptide.
+   - Dataset: 1 microsecond simulation of atom trajectories, mapped to coarse-grained states.
+   - Metrics: Relaxation times, stationary distributions.
 
-[More Information Needed]
+4. Simple Protein Folding (PFold):
 
-## Citation [optional]
+   - A 2-state model describing folding and unfolding rates of proteins.
+   - Dataset: Simulated transitions between folded and unfolded states.
+   - Metrics: Transition rates, mean first-passage times.
 
-<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
 
-**BibTeX:**
+#### Summary
 
-[More Information Needed]
+The Foundation Inference Model (FIM) represents a groundbreaking approach to zero-shot inference for Markov Jump Processes (MJPs). FIM enables accurate estimation of a variety of properties, including stationary distributions, relaxation times, mean first-passage times, time-dependent moments, and thermodynamic quantities (e.g., entropy production), all from noisy and discretely observed MJPs with state spaces of varying dimensionalities. Importantly, FIM operates in a zero-shot mode, requiring no additional fine-tuning or retraining on target datasets.
 
-**APA:**
+To the best of our knowledge, FIM is the first zero-shot generative model for MJPs, showcasing a versatile and powerful methodology for a wide range of physical, chemical, and biological systems. Future directions for FIM include extending its applicability to Birth and Death processes and incorporating more complex prior distributions for transition rates to enhance its generalization capabilities.
 
-[More Information Needed]
+## Limitations
 
-## Glossary [optional]
+While FIM has demonstrated strong performance on synthetic datasets, its methodology relies heavily on the distribution of these synthetic data. As a result, the model's effectiveness diminishes when evaluated on empirical datasets that significantly deviate from the synthetic distribution. For instance, as shown in Figure 4 (right), FIM's performance degrades rapidly for cases where the ratio between the largest and smallest transition rates exceeds three orders of magnitude. Such scenarios fall outside the range of FIM's prior Beta distributions and present challenges for accurate inference.
 
-<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+Additionally, the dynamics of MJPs underlying systems with long-lived, metastable states depend heavily on the shape of the energy landscape defining the state space. Transition rates in these systems are characterized by the depth of energy traps and can follow distributions not represented in FIM's training prior, such as power-law distributions (e.g., in glassy systems). These distributions lie outside the synthetic ensemble used for training FIM, limiting its ability to generalize to such cases.
 
-[More Information Needed]
+To address these limitations, future work will explore training FIM on synthetic MJPs with more complex transition rate distributions, such as those arising from systems with exponentially distributed energy traps or power-law-distributed rates, to better handle a broader range of real-world scenarios.
+## License
 
-## More Information [optional]
+The model is licensed under the Apache 2.0 License.
 
-[More Information Needed]
+## Citation
 
-## Model Card Authors [optional]
+If you use this model in your research, please cite:
 
-[More Information Needed]
+```
+@article{berghaus2024foundation,
+  title={Foundation Inference Models for Markov Jump Processes},
+  author={Berghaus, David and Cvejoski, Kostadin and Seifner, Patrick and Ojeda, Cesar and Sanchez, Ramses J},
+  journal={arXiv preprint arXiv:2406.06419},
+  year={2024}
+}
+```
 
-## Model Card Contact
+## Contact
 
-[More Information Needed]
+For questions or issues, please contact Kostadin Cvejoski at cvejoski@gmail.com.