# Sudoku noisy-latent trainer

This directory is **code only**. The training arrays are not in git.
They live on Hugging Face: https://huggingface.co/Avra98/Sudoku_superposition
under `data/`. They may already be on disk under another path.

If you are an agent: find the 11 files below (by **filename**, not folder name),
export the six path variables, then run `bash run_experiment.sh` from this directory.

## Required filenames

Search the machine for these exact names. A folder called `data/` is common but
not required.

| filename | role | env var |
|---|---|---|
| `train_sudoku_puzzles.npy` | raw train puzzles (clues + solution) | `SUDOKU_TRAIN_PATH` |
| `test_sudoku_puzzles.npy` | raw test puzzles | `SUDOKU_TEST_PATH` |
| `train_cand_masks.npy` | staged candidate sets, train `(N,12,81)` | `SUDOKU_TRAIN_CAND` |
| `test_cand_masks.npy` | staged candidate sets, test | `SUDOKU_TEST_CAND` |
| `train_assignments.npy` | train instance boards `(M,81)` | in `SUDOKU_INSTANCE_DIR` |
| `train_starts.npy` | train `(puzzle,stage)` row starts | in `SUDOKU_INSTANCE_DIR` |
| `train_counts.npy` | train `(puzzle,stage)` counts | in `SUDOKU_INSTANCE_DIR` |
| `test_assignments.npy` | test instance boards | in `SUDOKU_INSTANCE_DIR` |
| `test_starts.npy` | test starts | in `SUDOKU_INSTANCE_DIR` |
| `test_counts.npy` | test counts | in `SUDOKU_INSTANCE_DIR` |
| `test_bigq300.npz` | large-sample `q` for the hinge gate | `SUDOKU_GATE_Q` |

`SUDOKU_INSTANCE_DIR` is a **directory**. The loader opens
`{train,test}_{assignments,starts,counts}.npy` inside it.

Do **not** require `*_index.npy`, `*_meta.npy`, or `*_solution.npy`.

Typical places to look:

```
./data/
./Sudoku_superposition/data/
./datasets_superposition/          # rebuilt instances only
/tmp/sudoku_s12/                   # masks on some cluster nodes
/tmp/sudoku_superposition/         # instances on some cluster nodes
```

```bash
find "$HOME" /tmp /scratch -name 'train_assignments.npy' 2>/dev/null
find "$HOME" /tmp /scratch -name 'train_cand_masks.npy' 2>/dev/null
find "$HOME" /tmp /scratch -name 'train_sudoku_puzzles.npy' 2>/dev/null
```

## If the files are already on disk (any path)

Point each variable at the file you found. Masks and puzzles can live in one
folder and instances in another.

```bash
export SUDOKU_TRAIN_PATH=/abs/path/train_sudoku_puzzles.npy
export SUDOKU_TEST_PATH=/abs/path/test_sudoku_puzzles.npy
export SUDOKU_TRAIN_CAND=/abs/path/train_cand_masks.npy
export SUDOKU_TEST_CAND=/abs/path/test_cand_masks.npy
export SUDOKU_INSTANCE_DIR=/abs/dir/that/contains/assignments_starts_counts
export SUDOKU_GATE_Q=/abs/path/test_bigq300.npz

# this directory (the one that contains train/ and run_experiment.sh)
bash run_experiment.sh
```

Or, if all 11 files are in one folder:

```bash
SUDOKU_DATA_DIR=/abs/path/to/that/folder bash run_experiment.sh
```

`run_experiment.sh` refuses to start if any required file is missing or empty.

## If the files are not on disk

```bash
pip install -U "huggingface_hub[hf_xet]"
huggingface-cli download Avra98/Sudoku_superposition \
  --include "data/train_sudoku_puzzles.npy" \
  --include "data/test_sudoku_puzzles.npy" \
  --include "data/train_cand_masks.npy" \
  --include "data/test_cand_masks.npy" \
  --include "data/train_assignments.npy" \
  --include "data/train_starts.npy" \
  --include "data/train_counts.npy" \
  --include "data/test_assignments.npy" \
  --include "data/test_starts.npy" \
  --include "data/test_counts.npy" \
  --include "data/test_bigq300.npz" \
  --local-dir ./Sudoku_superposition

SUDOKU_DATA_DIR=$PWD/Sudoku_superposition/data bash run_experiment.sh
```

## What each group is

- **Puzzles** — clue tokens and solver order. Not replaceable by instances.
- **Masks** — post-constraint set `S` after each of 12 solver waves. Premade; keep them even if you rebuilt instances.
- **Instances** — concrete boards sampled from those masks, plus `starts`/`counts` for O(1) lookup.
- **Gate** — `test_bigq300.npz` for the hinge promotion statistic.

If instances were rebuilt, keep the premade masks and set `SUDOKU_INSTANCE_DIR` to the new folder only.

## Trainer command (what the script runs)

```bash
python -m train.main --workdir=../runs/noise --exp_name=noise
```

Must be run from this directory so `train/` is a package. Needs
`jax flax optax ml_collections clu absl-py tensorflow wandb`.

Useful knobs: `SUDOKU_LATENT_NOISE=1`, `SUDOKU_MC_SAMPLES=8`,
`SUDOKU_LATENT_SLOTS=12`, `SUDOKU_AUX_WEIGHT=0`, `SUDOKU_START_STAGE=1`
(set `4` to skip early curriculum).

Do not run `data_gen/staged_candidate_gen.py` or
`data_gen/build_superposition_dataset.py` just to start training.
Those already produced the HF `data/` files.