# Neural Characteristic Flow (NCF) for Classed-conditioned Optimal Transport

This repository contains the official implementation of our optimal transport model using neural networks and characteristics of Hamilton–Jacobi (HJ) equations. It focuses on the **class-conditioned OT tasks** discussed in Section 6.2 of the paper.

---

## Requirements

- All code is written and executed on **Google Colab** with **GPU acceleration** enabled.  
- Please upload the folder **`OTHJ_FashionMNIST_MNIST`** to your Colab environment before running the notebooks.

---

## Usage

### 1. Train the Classifier
Run all code blocks in **`train_classifier.ipynb`** to train the classifier on the MNIST dataset.

---

### 2. Pretrained VAEs

Run all code blocks in **`train_VAE_MNIST.ipynb`** or **`train_VAE_Fashion_MNIST`** to train the VAEs. For Fashion MNIST, we need to train one VAE with latent_space_dim=35; For MNIST, we need to train two VAEs with latent_space_dim=35 and 10, respectively.

After the excution of our codes, you should obtain the following directory structure:

```
OTHJ_FashionMNIST_MNIST/
├── example_d10_beta100_01/
│   └── best_model_dim_10_mnist.h5
└── example_d35_beta100_01/
    ├── best_model_dim_35_fmnist.h5
    └── best_model_dim_35_mnist.h5
```

---

### 3. Reproducing Section C.3 Tasks

Each task corresponds to a specific Jupyter notebook.  
Please open the notebook in Colab and **run all cells sequentially**.  
The main algorithm is implemented in the block *“OT HJ implicit solver”*.

#### Task 1: MNIST In-class Transfer  
Notebook: **`MNIST_in_class_transfer.ipynb`**

#### Task 2: MNIST In-class Shift  
Notebook: **`MNIST_in_class_shift.ipynb`**

#### Task 3: Fashion MNIST to MNIST Inter-class Transport
Notebook: **`FashionMNIST_MNIST.ipynb`**

---
