sparse_gpu_experiment

Can 16 point temperature sensors around a room detect that a person walked in — without ever seeing the full thermal field?

A sparse-sensor reconstruction model is trained only on ambient (AC-only) data from a GPU-accelerated 2D Boussinesq DNS simulation, then tested on a 150,000-step AC-only equilibration followed by a 50,000-step run where a mild-anomaly person enters, moves, and leaves. See reproduce_sparse_experiment.ipynb for the full walkthrough — what the model sees, exactly how its reconstruction error is computed, and the actual figures from both training runs below.

Stack

Layer	What / Where
CFD solver	Custom 2D Boussinesq finite-difference + projection-method DNS (solver_core_jax.py) — periodic-x via FFT, no-slip walls in z, Poisson solve via a hand-written parallel cyclic reduction (PCR) tridiagonal solver (not the sequential Thomas algorithm — see below)
GPU acceleration	JAX 0.11.0 (jax[cuda12]) + jax.jit / jax.lax.scan — fuses whole chunks of timesteps into single compiled GPU programs, no per-step Python dispatch
GPU runtime	WSL2 (Ubuntu, Python 3.14) — required because there is no CUDA-enabled jaxlib wheel for native Windows; GPU passthrough via the Windows host NVIDIA driver, no separate Linux driver needed
Hardware	NVIDIA GeForce RTX 3060 Ti (8GB), driver 591.86, CUDA 13.1
Numerics (WSL2 side)	NumPy 2.5.1, SciPy 1.18.0
ML model + training	PyTorch 2.5.1+cu121 — runs in the native Windows venv, not WSL2 (CUDA PyTorch has no such Windows-wheel gap)
Model architecture	MLP sensor-encoder (model_sparse.py, SparseToField) feeding the same conv-decoder as the project's main full-field ConvAutoencoder, for an apples-to-apples comparison
Analysis / visualization	NumPy 2.4.6, SciPy 1.17.1, Matplotlib 3.11.1 (SymLogNorm color scale, physical equal-aspect axes), scikit-learn 1.9.0 (PCA for latent-space analysis)
Video export	Matplotlib animation.PillowWriter (GIF, no ffmpeg dependency)
Notebook	Jupyter / nbconvert, executed end-to-end (no hidden state)

Why two Python environments: JAX's CUDA wheels only exist for Linux, so the DNS solver runs inside WSL2; PyTorch's CUDA wheels work fine on native Windows, so model training/inference/visualization stay there. Nothing in this folder needs both GPUs stacks running at once.

Two independent training runs, two models

• train_sparse_gpu.py → model_sparse_gpu.pt: trained on simulate_multivent_jax.py's 5 vent-position ambient datasets (domain randomization over vent location, for generalization — same idea as the main experiment's generalization_experiment/).
• train_sparse_longrun.py → model_sparse_longrun.pt: trained on simulate_test_transition_jax.py's single continuous 150k-step AC-only run (settled 80k-150k window only) — a single fixed vent position, but a much longer, better-equilibrated trajectory. Its held-out validation loss sits well above training loss — a real generalization gap, since a single long trajectory's tail isn't the same as genuine scenario diversity.

These are two different axes of "more data" (position diversity vs. sequence length/equilibration), not one superseding the other.

Pipeline

                 ┌─────────────────────────────┐
                 │  solver_core_jax.py (WSL2)   │
                 │  JAX + PCR Poisson solve     │
                 └──────────────┬──────────────┘
                                │
        ┌───────────────────────┼───────────────────────┐
        ▼                                                ▼
simulate_multivent_jax.py                 simulate_test_transition_jax.py
(5 vent positions, ambient only)          (150k AC-only -> 50k person enter/move/leave,
        │                                  checkpointed + resumable)
        ▼                                                ▼
train_sparse_gpu.py                       train_sparse_longrun.py
        │                                                │
        └───────────────────┬────────────────────────────┘
                             ▼
        visualize_sensor_patches.py / visualize_full_native.py
                             │
                             ▼
                     make_videos.py (GIFs)
                             │
                             ▼
              reproduce_sparse_experiment.ipynb (compiled writeup)

Regenerating data (needs WSL2 + the JAX venv, GPU)

python simulate_multivent_jax.py --positions 0.15,0.30,0.50,0.70,0.85 --max_dt 2e-4
python simulate_test_transition_jax.py --resume   # --resume picks up from
                                                    # checkpoint_transition.npz
                                                    # if a previous run was interrupted

simulate_test_transition_jax.py checkpoints every 5000 steps (atomic write) specifically because this machine's WSL2/GPU passthrough has been observed to stall mid-run — see the project's CLAUDE.md for the failure modes and fixes (smaller max_dt, a fixed timestep instead of recomputing every chunk, and a GPU-fused person-phase trajectory instead of small CPU-dispatched sub-chunks that left the GPU idle).

Training / evaluation (native Windows venv, PyTorch+CUDA — NOT WSL2)

python train_sparse_gpu.py            # or train_sparse_longrun.py
python visualize_sensor_patches.py --model model_sparse_longrun.pt --suffix _longrun
python visualize_full_native.py
python make_videos.py
jupyter nbconvert --to notebook --execute --inplace reproduce_sparse_experiment.ipynb

Data files (gitignored — regenerate, don't expect them in a fresh clone)

data/vent_x*.npy (5-position training set), data/baseline_quiet.npy / data/test_person.npy (150k/50k long-run train/test split), data/transition_full.npy + data/transition_meta.npz (64x64, whole long run), data/native_ambient.npy / data/native_person.npy (native-resolution snapshots for videos). All multi-hundred-MB to multi-GB — never commit these.