# iJKOnet: Learning Population Dynamics via Inverse JKO Optimization

## Overview
This repository provides an implementation of **iJKOnet**, an algorithm for **learning population dynamics** from temporal snapshots of particle distributions.  
The method combines the **Jordan–Kinderlehrer–Otto (JKO) scheme** with **inverse optimization** to recover the underlying **energy functional** that governs observed dynamics. Unlike prior JKO-based approaches, iJKOnet:

- Supports **end-to-end training** (no need for precomputed optimal transport maps).  
- Places **no architectural restrictions** (avoids input-convex neural networks).  
- Provides **theoretical guarantees** for energy recovery.  
- Demonstrates improved performance on synthetic and real-world datasets (e.g., single-cell genomics).  


---

## Method
- Models the **evolution of distributions** using the **Wasserstein gradient flow** framework.  
- Formulates the problem as a **min–max inverse optimization task** to learn the underlying energy functional.  
- Uses **Monte Carlo estimation** and adversarial training to optimize the loss.  
- Parameterizes:
  - Energy functional components: potential, interaction, diffusion.  
  - Transport maps: standard MLPs / ResNets with time encoding.  
