---
phase: 03-training-energy-tests
plan: 02
type: execute
wave: 2
depends_on: ["03-01"]
files_modified:
  - tests/test_energy.py
  - tests/test_trainer.py
  - src/pcn_torch/__init__.py
autonomous: true

must_haves:
  truths:
    - "Unit tests verify gain-modulated error uses preactivation f'(a^(l)) * eps not f'(x^(l)) * eps (TST-02)"
    - "Unit tests verify no autograd graph is built during training -- no RuntimeWarning emitted (TST-03)"
    - "Unit tests verify top_error = supervised_error @ W_out with correct shape (B, d_L) (TST-04)"
    - "Energy tests verify compute_energy returns a non-negative float from PCNErrors"
    - "Energy tests verify compute_energy_per_layer returns correct number of entries"
    - "Training integration tests verify train_pcn runs on classification and regression tasks"
    - "Training integration tests verify TrainHistory is populated with energy and accuracy data"
    - "Public API exports TrainConfig, TrainHistory, EnergyHistory, TrainCallback, RichCallback, train_pcn, test_pcn from pcn_torch"
  artifacts:
    - path: "tests/test_energy.py"
      provides: "Energy computation unit tests"
      min_lines: 40
    - path: "tests/test_trainer.py"
      provides: "Training loop correctness tests (TST-02, TST-03, TST-04) and integration tests"
      min_lines: 100
    - path: "src/pcn_torch/__init__.py"
      provides: "Public API surface for Phase 3 exports"
      contains: "train_pcn"
  key_links:
    - from: "tests/test_energy.py"
      to: "src/pcn_torch/energy.py"
      via: "from pcn_torch.energy import compute_energy, compute_energy_per_layer"
      pattern: "from pcn_torch\\.energy import"
    - from: "tests/test_trainer.py"
      to: "src/pcn_torch/trainer.py"
      via: "from pcn_torch.trainer import train_pcn, test_pcn, TrainConfig"
      pattern: "from pcn_torch\\.trainer import"
    - from: "src/pcn_torch/__init__.py"
      to: "src/pcn_torch/trainer.py"
      via: "re-exports all public trainer names"
      pattern: "from pcn_torch\\.trainer import"
---

<objective>
Write algorithmic correctness tests for PCN training (TST-02, TST-03, TST-04), energy computation tests, training integration tests, and update the public API exports.

Purpose: These tests verify structural correctness of the training algorithm without requiring numerical convergence. They catch the most dangerous bugs: wrong error formula, autograd leaks, and incorrect supervised error projection. The public API update makes all Phase 3 types importable from the top-level package.

Output: Two test files (test_energy.py, test_trainer.py) and updated __init__.py.
</objective>

<execution_context>
@C:\Users\RDCHLEM1\.claude/get-shit-done/workflows/execute-plan.md
@C:\Users\RDCHLEM1\.claude/get-shit-done/templates/summary.md
</execution_context>

<context>
@.planning/PROJECT.md
@.planning/ROADMAP.md
@.planning/STATE.md
@.planning/phases/03-training-energy-tests/03-RESEARCH.md
@.planning/phases/03-training-energy-tests/03-01-SUMMARY.md
@src/pcn_torch/energy.py
@src/pcn_torch/trainer.py
@src/pcn_torch/network.py
@src/pcn_torch/__init__.py
@tests/test_network.py
</context>

<tasks>

<task type="auto">
  <name>Task 1: Write energy and training correctness tests</name>
  <files>tests/test_energy.py, tests/test_trainer.py</files>
  <action>
All tests use synthetic data only: small deterministic tensors with `torch.manual_seed(42)`, dims=[20, 10, 5], batch_size=8, n_samples=16-32. Per user decision: structural tests only, no numerical convergence tests.

### tests/test_energy.py

**Helper:** Create a `_make_network_and_errors()` fixture function that:
- Sets torch.manual_seed(42)
- Creates PredictiveCodingNetwork(dims=[20, 10, 5])
- Calls init_latents(4)
- Creates x = torch.randn(4, 20), y = torch.randn(4, 5)
- Returns network, errors from compute_errors(x, y)

**Tests:**

1. `test_compute_energy_returns_float` -- verify return type is float (not Tensor)
2. `test_compute_energy_non_negative` -- energy >= 0
3. `test_compute_energy_without_supervised` -- compute_energy works when supervised_error is None (pass y=None to compute_errors)
4. `test_compute_energy_with_supervised` -- compute_energy includes supervised error term (energy with y > energy without y, since supervised error adds a positive term)
5. `test_compute_energy_manual_formula` -- manually compute (1/B) * sum(0.5 * ||e||^2) and compare to compute_energy output using torch.allclose on torch.tensor values
6. `test_compute_energy_per_layer_count` -- verify len(per_layer) == num_layers when no supervised, == num_layers + 1 when supervised
7. `test_compute_energy_per_layer_sum_matches_total` -- sum of per_layer entries approximately equals compute_energy total
8. `test_compute_energy_float16_safe` -- create errors with float16 tensors, verify compute_energy returns finite (not inf/nan)

### tests/test_trainer.py

**Shared fixture:** Create a `_make_dataloader()` helper:
```python
def _make_dataloader(n_samples=16, input_dim=20, num_classes=5, batch_size=8, seed=42):
    torch.manual_seed(seed)
    X = torch.randn(n_samples, input_dim)
    y = torch.randint(0, num_classes, (n_samples,))
    return DataLoader(TensorDataset(X, y), batch_size=batch_size, shuffle=False)
```

Also a `_make_regression_dataloader()` that returns continuous y values.

**TrainConfig tests:**

1. `test_config_defaults` -- verify default values match paper: T_infer=50, lr_infer=0.05, lr_learn=0.005, num_epochs=4
2. `test_config_invalid_task` -- ValueError for task="invalid"
3. `test_config_invalid_T_infer` -- ValueError for T_infer=0
4. `test_config_invalid_T_learn` -- ValueError for T_learn=0

**Algorithmic correctness tests (TST-02, TST-03, TST-04):**

5. `test_gain_modulated_error_uses_preactivation` (TST-02):
   - Create network dims=[20, 10, 5], init_latents(4), x = randn(4, 20)
   - errors = compute_errors(x)
   - Manual computation: layer0 = net.layers[0], states = [x, *net.latents]
   - preact = states[1] @ layer0.weight.T
   - eps = states[0] - layer0.activation_fn(preact)
   - expected_h = layer0.activation_deriv(preact) * eps
   - assert torch.allclose(errors.gm_errors[0], expected_h)

6. `test_no_autograd_graph_during_training` (TST-03):
   - Create network, dataloader, config with T_infer=3, T_learn=2, num_epochs=1
   - Run train_pcn inside warnings.catch_warnings(record=True)
   - Assert no RuntimeWarning containing "autograd" was emitted
   - After training, assert all latents have grad_fn is None and not requires_grad

7. `test_top_error_wout_projection` (TST-04):
   - Create network dims=[20, 10, 5], init_latents(4)
   - x = randn(4, 20), y = randn(4, 5)
   - errors = compute_errors(x, y)
   - Manual: y_hat = latents[-1] @ readout.weight.T, sup = y_hat - y, expected_top = sup @ readout.weight
   - assert torch.allclose(errors.top_error, expected_top)
   - assert errors.top_error.shape == (4, 5) -- shape is (B, d_L) not (B, d_out) (in this case both are 5 but the assertion documents intent)

**Training integration tests:**

8. `test_train_pcn_classification` -- train_pcn completes 1 epoch on classification data, history has populated energy.per_batch, energy.per_epoch, train_accuracy
9. `test_train_pcn_regression` -- train_pcn completes with task="regression", no accuracy tracked (train_accuracy is empty)
10. `test_train_pcn_energy_monotonic_per_batch` -- NOT convergence test; just verify per_batch list has expected length (num_batches * num_epochs)
11. `test_train_pcn_history_config_stored` -- history.config is the same config object that was passed in
12. `test_test_pcn_returns_metrics` -- test_pcn returns dict with "accuracy" and "energy" keys, both are floats
13. `test_train_pcn_weights_change` -- snapshot weights before training, verify they differ after training (weights actually updated)
14. `test_train_pcn_mixed_precision_false` -- runs without error when use_mixed_precision=False (CPU test)

**Use standard pytest patterns:** Each test is a standalone function (no class needed). Use `torch.manual_seed(42)` at the start of each test. Use `TrainConfig(T_infer=3, T_learn=2, num_epochs=1, use_mixed_precision=False)` for fast tests (structural, not convergence).

**Coding standards:**
- `from __future__ import annotations`
- Import from pcn_torch (public API) where possible, from submodules for internal testing
- No test classes needed -- flat functions are fine for this scale
  </action>
  <verify>
Run: `pytest tests/test_energy.py tests/test_trainer.py -v`
Run: `ruff check tests/test_energy.py tests/test_trainer.py`
All tests pass, no ruff errors.
  </verify>
  <done>
test_energy.py has 8 tests covering energy computation correctness including float16 safety. test_trainer.py has 14 tests covering TrainConfig validation, algorithmic correctness (TST-02, TST-03, TST-04), and training integration. All tests use synthetic data with seeded randomness. All pass.
  </done>
</task>

<task type="auto">
  <name>Task 2: Update public API exports in __init__.py</name>
  <files>src/pcn_torch/__init__.py</files>
  <action>
Update `src/pcn_torch/__init__.py` to export all Phase 3 public names:

Add imports:
```python
from pcn_torch.energy import compute_energy, compute_energy_per_layer
from pcn_torch.trainer import (
    EnergyHistory,
    RichCallback,
    TrainCallback,
    TrainConfig,
    TrainHistory,
    train_pcn,
    test_pcn,
)
```

Add to `__all__`:
```python
"compute_energy",
"compute_energy_per_layer",
"EnergyHistory",
"RichCallback",
"TrainCallback",
"TrainConfig",
"TrainHistory",
"train_pcn",
"test_pcn",
```

Maintain alphabetical order in __all__ (matching existing convention).
  </action>
  <verify>
Run: `python -c "from pcn_torch import train_pcn, test_pcn, TrainConfig, TrainHistory, EnergyHistory, compute_energy; print('OK')"`
Run: `ruff check src/pcn_torch/__init__.py`
Run: `mypy src/pcn_torch/__init__.py`

Full test suite:
Run: `pytest tests/ -v`
  </verify>
  <done>
All Phase 3 public names (train_pcn, test_pcn, TrainConfig, TrainHistory, EnergyHistory, TrainCallback, RichCallback, compute_energy, compute_energy_per_layer) are importable from pcn_torch. The full test suite passes.
  </done>
</task>

</tasks>

<verification>
1. `pytest tests/test_energy.py -v` -- all energy tests pass
2. `pytest tests/test_trainer.py -v` -- all trainer tests pass (including TST-02, TST-03, TST-04)
3. `pytest tests/ -v` -- full test suite passes (including Phase 1 and Phase 2 tests)
4. `ruff check src/pcn_torch/__init__.py tests/test_energy.py tests/test_trainer.py` -- no errors
5. `mypy src/pcn_torch/__init__.py` -- no errors
6. `python -c "from pcn_torch import train_pcn, test_pcn, TrainConfig"` -- succeeds
</verification>

<success_criteria>
- 8 energy computation tests pass (including float16 safety)
- 14 trainer tests pass (including TST-02 gain-modulated error, TST-03 no autograd, TST-04 W_out^T projection)
- All tests use synthetic data with seeded randomness (fast, deterministic, no downloads)
- Public API exports all Phase 3 types from pcn_torch top-level
- Full test suite (all phases) passes
- ruff and mypy pass on all modified files
</success_criteria>

<output>
After completion, create `.planning/phases/03-training-energy-tests/03-02-SUMMARY.md`
</output>
