iJKOnet

KO | EN

gitlite — search

#hacktoberfest #ai #typescript #llm #go #rust #linux #docker #java #cli #vue #security #mcp #ios

(★ 9)

[ICLR 2026] JAX implementation of the iJKOnet method

iJKOnet 최신버젼 다운로드

최종 버전 다운로드 (.zip)

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# Learning of Population Dynamics:<br> Inverse Optimization Meets JKO Scheme

[Mikhail Persiianov](https://scholar.google.com/citations?user=YXX3VMQAAAAJ&hl=en),
[Jiawei Chen](https://bluebluecard.github.io),
[Petr Mokrov](https://scholar.google.com/citations?user=CRsi4IkAAAAJ&hl=en),
[Alexander Tyurin](https://scholar.google.com/citations?user=Es8-xocAAAAJ&hl=en),
[Evgeny Burnaev](https://scholar.google.com/citations?user=pCRdcOwAAAAJ&hl=en),
[Alexander Korotin](https://scholar.google.com/citations?user=1rIIvjAAAAAJ&hl=en)

[![arXiv Paper](https://img.shields.io/badge/arXiv-2506.01502-b31b1b)](https://arxiv.org/abs/2506.01502)
[![OpenReview Paper](https://img.shields.io/badge/OpenReview-PDF-8c1b13)](https://openreview.net/forum?id=tVJIKd6CLF)
[![GitHub](https://img.shields.io/github/stars/MuXauJl11110/iJKOnet?style=social)](https://github.com/MuXauJl11110/iJKOnet)
![GitHub License](https://img.shields.io/github/license/MuXauJl11110/iJKOnet)

</div>

This repository contains the official implementation of the paper *"Learning of Population Dynamics: Inverse Optimization Meets JKO Scheme"*, accepted at **ICLR 2026**. The goal of this work is to learn *population-level dynamics from snapshot observations* using the theory of Wasserstein gradient flows and optimal transport.

## 📌 TL;DR

This paper introduces $\texttt{iJKOnet}$, a method for learning population dynamics from unpaired snapshot distributions. It recovers an energy functional governing a Wasserstein gradient flow via an inverse formulation of the JKO scheme, resulting in a min–max objective between the energy model and transport maps. The approach avoids inner JKO solvers and precomputed optimal transport plans, enabling end-to-end training with standard neural networks.

## 📦 Project Structure

```bash
.
├── configs/               # training configuration files
│   ├── config-base.yaml
│   └── config-method-*.yaml
│
├── data/                  # datasets
├── models/                # JKOnet, JKOnet*, iJKOnet implementations
├── networks/              # neural network architectures
│
├── notebooks/             # tutorial notebooks
│
├── scripts/               # experiment scripts
│   ├── bash/              # generated bash scripts
│   ├── generate_sc_w_lo.py
│   ├── generate_sc_wo_lo.py
│   └── optuna_search.sh
│
├── utils/                 # helper utilities
│   ├── dataset/
│   ├── entropy_estimation/
│   └── evaluation/
│
├── train.py               # main training entrypoint
└── optuna_search.py       # hyperparameter search
```

## 📥 Installation & Dependencies

Install the necessary packages for this repository by creating the Anaconda environment:

```bash
git clone https://github.com/.../iJKOnet.git
cd iJKOnet

conda env create -f environment.yml
conda activate ijkonet  # Replace with actual environment name if different
```

## **📔 Notebooks**

The `notebooks` folder contains example notebooks demonstrating how to generate data and train models. In particular, `notebooks/iJKOnet_usage.ipynb` demonstrates how to:

- generate synthetic 2D data for learning a potential energy function,

- train the available models on this data.

## **📜 Scripts**

The scripts

- `generate_sc_w_lo.py`
- `generate_sc_wo_lo.py`

generate bash scripts (saved in `scripts/bash`) for running
single-cell experiments:

• **leave-one-out experiments**  
• **full trajectory experiments (no leave-one-out)**

The script `scripts/optuna_search.sh` is used to launch a SLURM job for hyperparameter search.

## **🧬 Learning Single-Cell Dynamics**

For the [EB dataset](https://www.nature.com/articles/s41587-019-0336-3), we followed the preprocessing pipeline described in the [$\texttt{JKOnet}^\star$ tutorial](https://jkonet-star.readthedocs.io/en/latest/pages/tutorial_rna.html#id94).

For the [Multi dataset](https://openproblems.bio/events/2022-08_neurips), we used the preprocessing pipeline described in the paper [**“A Computational Framework for Solving Wasserstein Lagrangian Flows”**](https://arxiv.org/abs/2310.10649) and implemented in the corresponding [repository](https://github.com/necludov/wl-mechanics).

## 🏋️ Training

To start training, use the following general pattern:

```bash
python train.py \
  --solver <solver_name> \
  --dataset <dataset_name> \
  --config <path_to_base_config> \
  --extra_config <path_to_additional_config> \
  --K <K> \
  --array-tau <tau_or_tau_list> \
  --epochs <number_of_epochs> \
  --seed <seed>
```

## 🎓 Citation

If you find this repository useful in your research, please cite:

```bibtex
@inproceedings{
  persiianov2026learning,
  title={Learning of Population Dynamics: Inverse Optimization Meets {JKO} Scheme},
  author={Mikhail Persiianov and Jiawei Chen and Petr Mokrov and Alexander Tyurin and Evgeny Burnaev and Alexander Korotin},
  booktitle={The Fourteenth International Conference on Learning Representations},
  year={2026},
  url={https://openreview.net/forum?id=tVJIKd6CLF}
}
```

## 🙏 Credits

- [jkonet-star](https://github.com/antonioterpin/jkonet-star/tree/main) — our code is primarily based on this repository with some fixes of the data generation for synthetic experiments;
- [mutinfo](https://github.com/VanessB/mutinfo/tree/master) — the code in `utils/entropy_estimation` is based on code from this repository;
- [POT](https://pythonot.github.io) and [ott-jax](https://ott-jax.readthedocs.io) for ot toolkit;
- [optuna](https://optuna.org) - toolkit for hyperparameter search
- [comet ML](https://www.comet.com) — experiment-tracking and visualization toolkit;
- [inkscape](https://inkscape.org/) — an excellent open-source editor for vector graphics;

iJKOnet

iJKOnet 최신버젼 다운로드

오늘의 추천 프로젝트

Was this content helpful?

Recent Feedback

iJKOnet

iJKOnet 최신버젼 다운로드

오늘의 추천 프로젝트

Was this content helpful?

Recent Feedback

Download README