| import numpy as np |
| import torch |
|
|
| from convgru_ensemble.lightning_model import RadarLightningModel |
|
|
|
|
| def test_predict_handles_unpadded_inputs(): |
| model = RadarLightningModel( |
| input_channels=1, |
| num_blocks=1, |
| forecast_steps=2, |
| ensemble_size=1, |
| noisy_decoder=False, |
| ) |
| past = np.zeros((4, 8, 8), dtype=np.float32) |
|
|
| preds = model.predict(past, forecast_steps=2, ensemble_size=1) |
|
|
| assert preds.shape == (1, 2, 8, 8) |
| assert np.isfinite(preds).all() |
|
|
|
|
| def test_from_checkpoint_delegates_to_lightning_loader(monkeypatch): |
| captured = {} |
|
|
| def fake_loader(cls, checkpoint_path, map_location=None, strict=None, weights_only=None): |
| captured["checkpoint_path"] = checkpoint_path |
| captured["map_location"] = map_location |
| captured["strict"] = strict |
| captured["weights_only"] = weights_only |
| return "loaded-model" |
|
|
| monkeypatch.setattr(RadarLightningModel, "load_from_checkpoint", classmethod(fake_loader)) |
|
|
| loaded = RadarLightningModel.from_checkpoint("/tmp/model.ckpt", device="cpu") |
|
|
| assert loaded == "loaded-model" |
| assert captured["checkpoint_path"] == "/tmp/model.ckpt" |
| assert isinstance(captured["map_location"], torch.device) |
| assert captured["map_location"].type == "cpu" |
| assert captured["strict"] is True |
|
|
