source/gconvex_authors/scripts/run_finite_ot.py

"""Script to run FCOT-Separable on DGPS data and plot running statistics."""

import logging
import os

import numpy as np
import torch

from optimal_transport.ot_fc_sep_map import FCOTSeparable
from tools.dgps import generate_gaussian_pairs, generate_grid_XY
from tools.feedback import set_log_level
from tools.utils import (
    L22_1d,
    L33_1d,
    L44_1d,
    L55_1d,
    inverse_L22x,
    inverse_L33x,
    inverse_L44x,
    inverse_L55x,
    nL22_1d,
    inverse_nL22x,
)


SEED = 42
LOG_LEVEL = "DEBUG"

# Data distribution settings (1D Gaussian → Gaussian)
N_SAMPLES = 10000
MU_X_1D = torch.tensor([0.0])
SIGMA_X_1D = torch.tensor([[1.0]])
MU_Y_1D = torch.tensor([-1.0])
SIGMA_Y_1D = torch.tensor([[2.0]])

# Domain / grid resolution (matches target grid L in data generation)
RADIUS = 4.0
X_ACCURACY = 2e-3
Y_ACCURACY = 2e-3

# Optimizer (intercept optimizer)
OUTER_LR = 1e-2

# Softmax temperature schedule
TEMP_MIN = 1.0
TEMP_MAX = 60.0
TEMP_WARMUP_ITERS = 2_500

# Reactivation and full refresh rules
REACTIVATE_EVERY = 50
REACTIVATE_EPS = 1e-3
FULL_REFRESH_EVERY = 300

# Coarse-to-fine transform speedup
COARSE_X_FACTOR = 100
COARSE_TOP_K = 4
COARSE_WINDOW = 1

# Training / evaluation parameters
TRAIN_ITERS = 7_000
PRINT_EVERY = 10
LOG_EVERY = 20
CONVERGENCE_TOL = 1e-6
CONVERGENCE_PATIENCE = 300
FORCE_RETRAIN = False

# 1D evaluation settings
KERNELS_1D = [L22_1d, nL22_1d]
INV_KERNELS = [inverse_L22x, inverse_nL22x]


def _setup_randomness_and_threads():
    """Seed random generators and limit thread usage for reproducible runs."""
    torch.manual_seed(SEED)
    set_log_level(LOG_LEVEL)

    cpu = max(1, (os.cpu_count() or 1) // 2)
    try:
        torch.set_num_threads(cpu)
        torch.set_num_interop_threads(max(1, cpu // 2))
    except RuntimeError:
        pass


def _build_solver(
    dim: int,
    kernel_1d,
    inverse_kx,
    *,
    radius_override: float | None = None,
    x_accuracy_override: float | None = None,
    y_accuracy_override: float | None = None,
    ny_override: int | None = None,
) -> FCOTSeparable:
    radius = RADIUS if radius_override is None else radius_override
    x_accuracy = X_ACCURACY if x_accuracy_override is None else x_accuracy_override
    y_accuracy = Y_ACCURACY if y_accuracy_override is None else y_accuracy_override
    if ny_override is None:
        ny = int((2 * radius) / y_accuracy) + 1
    else:
        if ny_override < 2:
            raise ValueError("ny_override must be at least 2 grid points.")
        ny = ny_override
    n_params = ny * dim

    """Instantiate FCOTSeparable solver with the requested kernel/grid resolution."""
    fcot = FCOTSeparable.initialize_right_architecture(
        dim=dim,
        radius=radius,
        n_params=n_params,
        x_accuracy=x_accuracy,
        kernel_1d=kernel_1d,
        inverse_kx=inverse_kx,
        outer_lr=OUTER_LR,
        temp_min=TEMP_MIN,
        temp_max=TEMP_MAX,
        temp_warmup_iters=TEMP_WARMUP_ITERS,
        reactivate_every=REACTIVATE_EVERY,
        reactivate_eps=REACTIVATE_EPS,
        full_refresh_every=FULL_REFRESH_EVERY,
        cache_gradients=True,
        coarse_x_factor=COARSE_X_FACTOR,
        coarse_top_k=COARSE_TOP_K,
        coarse_window=COARSE_WINDOW,
    )
    return fcot


def _compute_transport(solver: FCOTSeparable, X: torch.Tensor) -> torch.Tensor:
    """Helper to evaluate the learned Monge map (c-gradient) on new points."""
    X_req = X.to(solver.device).requires_grad_(True)
    _, u_X = solver.model.forward(X_req, selection_mode="hard")
    grad_u = torch.autograd.grad(u_X.sum(), X_req, create_graph=False)[0]
    Y_pred = solver.inverse_kx(X_req.detach(), grad_u.detach())
    return Y_pred


def run_gaussian_1d():
    """Run 1D Gaussian→Gaussian OT experiments for multiple kernels."""
    _setup_randomness_and_threads()

    params = {
        "n": N_SAMPLES,
        "μ_x": MU_X_1D,
        "Σ_x": SIGMA_X_1D,
        "μ_y": MU_Y_1D,
        "Σ_y": SIGMA_Y_1D,
    }
    X, Y, _ = generate_gaussian_pairs(**params)

    sx = SIGMA_X_1D[0, 0].item()
    sy = SIGMA_Y_1D[0, 0].item()
    sigma_ratio = np.sqrt(sy / sx)

    def T_theory(x):
        return sigma_ratio * x - 1.0

    for kernel_1d, inverse_kx in zip(KERNELS_1D, INV_KERNELS):
        name = getattr(kernel_1d, "__name__", str(kernel_1d))
        print(f"\n[RUN 1D] Using kernel: {name}")

        fcot = _build_solver(dim=1, kernel_1d=kernel_1d, inverse_kx=inverse_kx)
        fcot.fit(
            X, Y,
            iters=TRAIN_ITERS,
            print_every=PRINT_EVERY,
            log_every=LOG_EVERY,
            convergence_tol=CONVERGENCE_TOL,
            convergence_patience=CONVERGENCE_PATIENCE,
            force_retrain=FORCE_RETRAIN,
        )

        X_test = torch.linspace(-3, 3, TEST_POINTS_1D).reshape(-1, 1)
        Y_pred = _compute_transport(fcot, X_test)
        Y_true = T_theory(X_test)

        mae = (Y_pred - Y_true).abs().mean().item()
        print(f"[RUN 1D] kernel={name} MAE vs theory = {mae:.4f}")


def run_gaussian_2d_grid():
    """Run 2D grid-like OT experiments for multiple kernels."""
    _setup_randomness_and_threads()

    centers = [-3.0, 0.0, 3.0]
    X, Y = generate_grid_XY(
        n=N_SAMPLES,
        L=RADIUS,
        std=0.45,
        centers=centers,
        force=True,
    )
    X = X.clamp(-RADIUS + 1e-3, RADIUS - 1e-3)
    Y = Y.clamp(-RADIUS + 1e-3, RADIUS - 1e-3)

    for kernel_1d, inverse_kx in zip(KERNELS_1D, INV_KERNELS):
        name = getattr(kernel_1d, "__name__", str(kernel_1d))
        print(f"\n[RUN 2D] Using kernel: {name}")

        fcot = _build_solver(dim=2, kernel_1d=kernel_1d, inverse_kx=inverse_kx)
        fcot.fit(
            X, Y,
            iters=TRAIN_ITERS,
            print_every=PRINT_EVERY,
            log_every=LOG_EVERY,
            convergence_tol=CONVERGENCE_TOL,
            convergence_patience=CONVERGENCE_PATIENCE,
            force_retrain=FORCE_RETRAIN,
        )

        Y_pred = _compute_transport(fcot, X)

        mean_X = X.mean(dim=0)
        mean_Y = Y.mean(dim=0)
        mean_Ypred = Y_pred.mean(dim=0)

        dist_initial = (mean_X - mean_Y).norm().item()
        dist_mapped = (mean_Ypred - mean_Y).norm().item()
        print(
            f"[RUN 2D] kernel={name} "
            f"mean dist initial={dist_initial:.4f}, mapped={dist_mapped:.4f}"
        )


if __name__ == "__main__":
    # run_gaussian_1d()
    run_gaussian_2d_grid()