File size: 3,898 Bytes
ae73c7f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | """
End-to-end test of run_multistream.py with the_well.data.WellDataset
mocked at the point env.py imports it. This exercises the FULL real path
(streaming -> adapter -> per-domain normalizer -> channel-filtered
replay -> training) together.
"""
from __future__ import annotations
import os
import sys
from unittest.mock import patch, MagicMock
import pytest
import torch
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from src.model import MultiScaleEncoder, HierarchicalHyperbolicPredictor
from src.continual import ReplayBuffer, DiagonalEWC, fit_normalizer_for_domain
from src.config import BEST_HPARAMS as BEST
from src.provenance import DataLoadError
def make_fake_well_dataset(n_samples, Ti, To, H, W, C):
class FakeWellDataset:
def __init__(self, **kwargs):
self._n = n_samples
def __len__(self):
return self._n
def __getitem__(self, idx):
return {
"input_fields": torch.randn(Ti, H, W, C),
"output_fields": torch.randn(To, H, W, C),
}
return FakeWellDataset
def test_load_real_domain_hard_fails_when_stream_unavailable():
from src.run_multistream import load_real_domain
with patch("the_well.data.WellDataset", side_effect=ConnectionError("no network")):
with pytest.raises(DataLoadError) as exc_info:
load_real_domain("nonexistent_dataset", "train", max_samples=8)
assert exc_info.value.outcome_code == "STREAM_FAILED"
def test_load_real_domain_succeeds_with_mocked_stream():
from src.run_multistream import load_real_domain
fake_cls = make_fake_well_dataset(n_samples=16, Ti=4, To=4, H=16, W=16, C=2)
with patch("the_well.data.WellDataset", fake_cls):
ds = load_real_domain("fake_dataset", "train", max_samples=16)
assert ds.provenance == "REAL_STREAMED"
item = ds[0]
assert item["fields"].shape == (8, 2, 16, 16) # Ti+To, C, H, W
def test_full_multistream_pipeline_two_domains_different_channels():
from src.run_multistream import load_real_domain, train_one_domain
fake_c2 = make_fake_well_dataset(n_samples=16, Ti=4, To=4, H=16, W=16, C=2)
fake_c11 = make_fake_well_dataset(n_samples=16, Ti=4, To=4, H=16, W=16, C=11)
enc = MultiScaleEncoder(hidden=32, out_dim=8)
model = HierarchicalHyperbolicPredictor(enc, c=1.0, pred_steps=2, levels=2)
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
replay = ReplayBuffer(capacity=64)
ewc = DiagonalEWC(model, lambda_ewc=100.0)
with patch("the_well.data.WellDataset", fake_c2):
ds1 = load_real_domain("gray_scott_reaction_diffusion", "train", max_samples=16)
assert ds1.provenance == "REAL_STREAMED"
norm1 = fit_normalizer_for_domain(ds1, max_fit=8)
before = model.encoder.stem.weight.detach().clone()
train_one_domain(model, norm1, ds1, opt, "cpu", epochs=1,
replay=replay, ewc=None, teacher=None, mix_replay=0.0)
after_domain1 = model.encoder.stem.weight.detach().clone()
assert not torch.allclose(before, after_domain1)
assert replay.counts_by_channels().get(2, 0) > 0
with patch("the_well.data.WellDataset", fake_c11):
ds2 = load_real_domain("active_matter", "train", max_samples=16)
assert ds2.provenance == "REAL_STREAMED"
norm2 = fit_normalizer_for_domain(ds2, max_fit=8)
train_one_domain(model, norm2, ds2, opt, "cpu", epochs=1,
replay=replay, ewc=ewc, teacher=None, mix_replay=0.5)
after_domain2 = model.encoder.stem.weight.detach().clone()
assert not torch.allclose(after_domain1, after_domain2)
assert replay.counts_by_channels().get(11, 0) > 0
assert set(replay.counts_by_channels().keys()) == {2, 11}
assert norm1.mean.shape == (2,)
assert norm2.mean.shape == (11,)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v"]))
|