# This file is part of the paper:
#
#    "Overcoming slow Kolmogorov width decay in parametric optimal control
#     via neural network surrogates"
#
#    https://github.com/HenKlei/PDOCP.git
#
# Copyright 2026 all developers. All rights reserved.
# License: BSD-2-Clause (http://opensource.org/licenses/BSD-2-Clause)
# Authors: Hendrik Kleikamp, Martin Lazar, Juan Ricardo Muñoz

Overcoming slow Kolmogorov width decay in parametric optimal control via neural network surrogates

In this repository, we provide the code used for the numerical experiments in the paper "Overcoming slow Kolmogorov width decay in parametric optimal control via neural network surrogates" by Hendrik Kleikamp, Martin Lazar and Juan Ricardo Muñoz.

Installation

Editable install from the project root:

python -m venv venv
. venv/bin/activate
pip install -e .

This makes pdocp importable from any working directory and pulls the four runtime dependencies (numpy, scipy, torch, matplotlib). Minimum versions are specified in pyproject.toml.

Optional: GPU support

pip install -e . installs the default PyTorch wheel, which is CPU-only on Linux. For CUDA support, first install a matching torch build following the selector at https://pytorch.org/get-started/locally, then run pip install -e ..

Reproducing the paper

Each paper example is a single training-set-size sweep across the methods that appear in the paper.

Example 1: moving rectangle

python sweep.py --setup rect_full \
    --methods unet greedy param_dec cnn_ae_rom \
    --seeds 42 43 44 --n-restarts 3 \
    --outdir results/rect_full/sweep \
    --device cuda

LaTeX tables and pgfplots .dat files:

python make_latex_outputs.py --setup rect_full \
    --results-dir results/rect_full/sweep \
    --outdir results/rect_full/figures

Example 2: moving ball with moving interface

python sweep.py --setup ball_interface_joint_square \
    --methods unet greedy param_dec cnn_ae_rom \
    --seeds 42 43 44 --n-restarts 3 \
    --outdir results/ball_interface_joint_square/sweep \
    --device cuda

LaTeX tables and .dat files:

python make_latex_outputs.py --setup ball_interface_joint_square \
    --results-dir results/ball_interface_joint_square/sweep \
    --outdir results/ball_interface_joint_square/figures

Presentation plots

Two scripts produce the per-quantity, borderless PNG+PDF panels used in the paper figures (one file per field, suitable for inclusion via tikz/pgfplots).

Input fields and FOM final time adjoints (no surrogate)

python make_presentation_plots.py --setup ball_interface_joint_square

U-Net predictions vs FOM ground truth

Loads a saved U-Net checkpoint, runs it on chosen test parameters, and emits one borderless image per quantity per parameter (target, FOM phi_T, U-Net phi_T, pointwise error, recovered x(T), plus the diffusivity field where it varies). A control-domain outline is overlaid automatically on setups with restricted control.

Ball setup (best n_train=1000 cell was seed 44; three test parameters that span the geometry: ball fully inside the control square, ball overhanging the boundary, ball intersecting the diffusivity interface):

python plot_unet_predictions.py \
    --setup ball_interface_joint_square \
    --weights 'results/ball_interface_joint_square/sweep/ball_interface_joint_square_unet_n1000_s44_UNet[xT+kappa]_weights.pt' \
    --snapshots results/ball_interface_joint_square/sweep/ball_interface_joint_square_snapshots_N64_M1000_bdf2.npz \
    --unet-channels xT kappa \
    --n-train 1000 --seed-split 44 --m-weight 1000 \
    --mu-indices 873 718 208 \
    --out-prefix results/ball_interface_joint_square/figures/

Each parameter writes one _info.txt recording the parameter values, per-sample phi_rel, and the colour-bar ranges used.

Questions

If you have any questions, feel free to contact us via email at juan.munoz@dim.uchile.cl.

License (BSD-2-Clause)

The code is licensed under BSD-2-clause, see LICENSE.txt.