Extropic/THRML optimization: Max-Cut

This repository was built for my article "How Extropic's Thermodynamic Computer Works". It contains two runnable scripts that demonstrate a simple optimization problem (Max‑Cut) implemented:

• naively (simulated annealing), and
• with THRML (a JAX library that emulates thermodynamic block Gibbs sampling).

What problem does this solve?

We need to split n employees into two shifts (Day/Night) to minimize conflicts (maximize cut). Each edge weight encodes how costly it is to put two employees on the same shift (think "how much A dislikes working with B"). A cut assigns employees to Day/Night; the cut weight is the sum of weights on edges that cross between the shifts. Higher cut weight means more high‑conflict pairs are separated, i.e. fewer conflicts remain within any single shift. If W_total is the sum of all edge weights, then "conflict left within shifts" is W_total − cut_weight.

• src/naive_max_cut.py: simulated annealing baseline.
• src/thrml_max_cut.py: THRML Ising Gibbs sampler.

Setup

uv venv
source .venv/bin/activate
pip install -r requirements.txt

Notes:

• JAX on CPU/macOS is fine for small demos but slower; a GPU/Metal build speeds up THRML.

Usage

Quality metric: higher "cut weight" is better (sum of weights on edges crossing Day/Night, it reflects separating high‑conflict pairs).

Parameters (both scripts):

• --n: number of nodes.
• --p: G(n,p) edge probability.
• --wmin/--wmax: weight range (uniform).
• --seed: RNG seed for reproducibility.

Naive SA specific:

• --steps: annealing sweeps.
• --beta-start/--beta-end: inverse‑temperature schedule.

THRML specific:

• --beta: fixed inverse temperature for Gibbs.
• --warmup: sweeps before sampling.
• --samples: number of saved samples.
• --steps-per-sample: sweeps between samples.
• --compile-separate: exclude JIT compile from timing by priming a tiny run.

Apples‑to‑apples comparison

Match total sweeps: sweeps_naive = steps; sweeps_thrml = warmup + (samples−1)×steps_per_sample.

Example (same graph/seed, 300 sweeps, fixed β):

# Naive (β fixed)
python src/naive_max_cut.py --n 80 --p 0.2 --steps 300 --beta-start 8.0 --beta-end 8.0 --seed 0

# THRML (β fixed, 100 + (101−1)*2 = 300 sweeps), exclude JIT time
python src/thrml_max_cut.py --n 80 --p 0.2 --beta 8.0 --warmup 100 --samples 101 --steps-per-sample 2 --seed 0 --compile-separate

Acceleration

By default, runs on CPU. To use an accelerator:

• Apple Silicon (Metal): pip install -U "jax>=0.4.30" "jax-metal"
• Linux + NVIDIA CUDA: pip install -U "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

THRML is accelerator‑friendly; larger graphs benefit more. See the overview: THRML examples: probabilistic computing.

How the THRML example works (high‑level)

• We map Max‑Cut to an Ising EBM (anti‑ferromagnetic couplings), build an IsingEBM over SpinNodes, run block Gibbs, and report the best cut weight found.

Official resources

• THRML documentation: https://docs.thrml.ai/en/latest/
• Inside X0 and XTR‑0: https://extropic.ai/writing/inside-x0-and-xtr-0