| """ |
| 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) |
|
|
|
|
| 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"])) |
|
|
