RGA (1D, b-grid) for 1D Poisson Problem

This repository implements the Relaxed Greedy Algorithm (RGA) to solve a 1D Poisson equation on the interval [0, \pi]. The implementation uses JAX for JIT compilation and GPU acceleration.

The script solves the variational problem u'' = -f with f = 1 and Robin-type boundary conditions enforced via a penalty term \frac{1}{2\delta}(u(0)^2 + u(\pi)^2) added to the energy. The exact solution of the underlying Dirichlet problem (u(0)=u(\pi)=0) is u^*(x) = 0.5 , x , (\pi - x).

Theoretical Background

The algorithm used here is based on the work of Siegel et al. (2023). It solves a minimization problem over the convex hull \mathcal{K}_1(\mathbb{D}) of a symmetric dictionary \mathbb{D}.

Relaxed Greedy Algorithm (RGA)

RGA is a greedy method for minimizing a convex and smooth objective \mathcal{R}: H \to \mathbb{R} over the convex hull of a dictionary \mathbb{D}.

Algorithm (RGA):

1. Initialization: u_0 = 0.
2. For n = 1 to N repeat: a. Compute the gradient: r = \nabla \mathcal{R}(u_{n-1}). b. Find the best atom (direction) in the dictionary \mathbb{D}: g^\star \in \arg\max_{g \in \mathbb{D}} |\langle g, r\rangle_H| c. Set the signed direction: g_n = \mathrm{sign}(\langle g^\star, r\rangle_H) , g^\star d. Set the relaxation step size: \alpha_n = \min(1, 2/n) e. Perform the update (with budget M): u_n = (1-\alpha_n)u_{n-1} - M , \alpha_n , g_n
3. Return: u_N.

Convergence rate and error analysis

The total error analysis splits into three main components:

1. Optimization error: RGA guarantees a convergence rate of \mathcal{O}(1/n) for the optimization error, where n is the number of iterations:

   \mathcal{R}(u_n) - \inf_{|v|_{\mathcal{K}_1(\mathbb{D})} \le M} \mathcal{R}(v) \le \frac{32 C M^2}{n}

2. Discretization error: The error introduced by numerical quadrature of the integrals (e.g. the trapezoidal rule on a fine grid). With high-order quadrature (as used here, trapezoid rule on a fine grid), this error is \mathcal{O}(N^{-(k+1)/d}), where k is the smoothness and d is the dimension.
3. Modeling error: The error that measures how well the "true" solution u^* can be approximated by functions in the convex hull of the dictionary. For Barron-type functions this component is typically smaller than the optimization error.

An \mathcal{O}(1/n) decay of the energy gap (as above) implies, via a Poincaré-type inequality, an \mathcal{O}(n^{-1/2}) decay in the $L^2$-norm:

|u_n - u^\star|*{L^2} \le C , \sqrt{\mathcal{R}(u_n) - \mathcal{R}(u^\star)} \implies |u_n - u^\star|*{L^2} = \mathcal{O}(n^{-1/2})

Implementation Details

The script main.py implements this algorithm with the following features:

• Technology: All numerical work is done in JAX, enabling JIT compilation (@jax.jit) and seamless CPU/GPU execution.

• Dictionary \mathbb{D}: The atoms are g(x) = S(x) \cdot \phi(wx + b), where:

  • \phi is a selectable activation function (e.g. sin, relu, relu2, relu3, tanh).
  • S(x) is an optional shape/envelope function (e.g. sin, xpix for x(\pi-x)) that can help atoms respect boundary behavior.

• argmax subproblem: The \arg\max over the dictionary parameters (w, b) is solved numerically:

  • The frequencies w are chosen from a fixed user-provided list (--w_values).
  • The bias b is searched over a predefined fine grid (the "b-grid") in the range [-\pi, \pi] (--Nb).
  • The scan over the b-grid is done in memory-friendly "chunks" of size --b_chunk.

• Update variants: The script implements two update strategies:

  1. Standard RGA (no --line_search): uses the theoretical relaxation step size \alpha_n = \min(1, 2/n).
  2. Line Search (with --line_search): computes an "optimal" step size \lambda = \frac{\langle r, g \rangle}{|g|_H^2} to minimize the energy in the chosen direction g.

• Options:

  • --norm_atoms: Normalizes each atom g_n before using it in the update.
  • --clip_budget: (Line-search only) Ensures that the accumulated variation k_1 respects the budget M by clipping the step size \lambda.

Setup and Execution

1. Virtual Environment (venv)

It is recommended to use a virtual environment.

# Create a virtual environment
python -m venv venv

# Activate it
# On macOS/Linux:
source venv/bin/activate
# On Windows:
# .\venv\Scripts\activate

2. Installation

Install the required packages from requirements.txt.

pip install -r requirements.txt

3. Execution

The main.py script can be called directly with arguments. The run.sh script can be used as a template.

Example run.sh:

#!/usr/bin/env bash
# run_paper.sh — runs all paper experiments and stores plots in a structured layout
set -euo pipefail

# Path to main.py (override if needed: MAIN=/path/to/main.py ./run_paper.sh)
MAIN="${MAIN:-main.py}"

# Base output directory (optionally provided as first argument)
OUT_ROOT="${1:-out/paper_runs}"

# Timestamped run directory
STAMP="$(date +%Y%m%d_%H%M%S)"
OUT="${OUT_ROOT}/${STAMP}"

# Shared parameters
ITERS=1000000
COMMON_ARGS=(--N 2001 --Nb 1200 --w_values=-1,1 --b_chunk 256 --log_every 5000)

# Checks
[[ -f "$MAIN" ]] || { echo "ERROR: $MAIN not found."; exit 1; }

mkdir -p "$OUT"

# Helper: run an experiment, create folder, save PNG+SVG+LOG
run_save () {
  local save_prefix="$1"; shift
  local dir
  dir="$(dirname "$save_prefix")"
  mkdir -p "$dir"

  local png="${save_prefix}.png"
  local svg="${save_prefix}.svg"
  local log="${save_prefix}.log"

  if [[ -f "$png" && -f "$svg" ]]; then
    echo "Skipping (already exists): $save_prefix"
    return 0
  fi

  echo ">>> Start: $save_prefix"
  python "$MAIN" "${COMMON_ARGS[@]}" --iters "$ITERS" "$@" --save "$save_prefix" |& tee "$log"
  echo ">>> Done:  $save_prefix"
}

echo "Output directory: $OUT"
echo "------------------------------------------------------------"

# 1) relu^2 (unnormalized), delta=1e-2, M=10
BASE1="$OUT/01_relu2_baseline"
run_save "$BASE1/relu2_d1e-2_M10" \
  --activation relu2 --delta 1e-2 --M 10 --shape sin2

# 2) relu2 normalized, M = 10,25,50,75,100; delta=1e-2
BASE2="$OUT/02_relu2_norm"
for M in 10 25 50 75 100; do
  run_save "$BASE2/relu2_norm_M${M}_d1e-2" \
    --activation relu2 --delta 1e-2 --norm_atoms --M "$M" --shape sin2
done

# 3) All shape functions UNNORMALIZED, delta=1e-2, M=10
BASE3="$OUT/03_shapes_unnorm"
for SH in none sin hump x2xpi xpix sin2; do
  run_save "$BASE3/relu2_shape-${SH}_unnorm_M10_d1e-2" \
    --activation relu2 --delta 1e-2 --M 10 --shape "$SH"
done

# 4) Line search with clipping (M=1000), deltas: 1e-2, 1e-4, 1e-5
BASE4="$OUT/04_linesearch_clip_M1000"
for DEL in 1e-2 1e-4 1e-5; do
  run_save "$BASE4/relu2_LSclip_M1000_d${DEL}" \
    --activation relu2 --delta "$DEL" --M 1000 --shape sin2 \
    --line_search --clip_budget
done
BASE4="$OUT/04_linesearch_clip_M10000"
for DEL in 1e-2 1e-3 1e-4; do
  run_save "$BASE4/relu2_LSclip_M10000_d${DEL}" \
    --activation relu2 --delta "$DEL" --M 10000 --shape sin2 \
    --line_search --clip_budget
done





echo "------------------------------------------------------------"
echo "All runs completed. Results (PNG, SVG, LOG) are stored under:"
echo "  $OUT"

Make executable and run:

chmod +x run.sh
./run.sh

Arguments

Key arguments for main.py:

• --save: Prefix for output files (plots).
• --iters: Total number of iterations.
• --M: Budget parameter M.
• --N: Number of grid points for x in the interval [0, \pi].
• --Nb: Number of grid points for b in the interval [-\pi, \pi].
• --activation: Activation function (sin, relu, relu2, relu3, tanh).
• --w_values: Comma-separated list of frequencies w (e.g. "1,2,3").
• --shape: Envelope/shape function (none, sin, xpix, sin2).
• --line_search: Enables the line-search update rule.
• --norm_atoms: Normalizes the atoms.
• --clip_budget: Enables budget clipping (line-search only).

Repository Structure

.
├── .gitignore
├── main.py           # Main script
├── readme.md         # This file
├── requirements.txt  # Python dependencies
├── run.sh            # Example run script
└── out/              # Directory for generated plots (in .gitignore)

Reference / Citation

The theoretical foundations of the RGA algorithm come from:

@article{Siegel2023Greedy,
  title = {Greedy training algorithms for neural networks and applications to PDEs},
  journal = {Journal of Computational Physics},
  volume = {484},
  pages = {112084},
  year = {2023},
  issn = {0021-9991},
  doi = {https://doi.org/10.1016/j.jcp.2023.112084},
  url = {https://www.sciencedirect.com/science/article/pii/S0021999123001791},
  author = {Jonathan W. Siegel and Qingguo Hong and Xianlin Jin and Wenrui Hao and Jinchao