File size: 6,136 Bytes
4ebcd24 | 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 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | """Regression checks for the capture-only released 1.2 pressure export."""
import hashlib
import json
import sys
import tempfile
import unittest
from pathlib import Path
import numpy as np
import torch
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "models/trackformer_1_2_field"))
from predict import export_pressure_fields
from plot_pressure import render_pressure_map
class PressureFieldExportTest(unittest.TestCase):
def setUp(self):
self.contract = {"normalization": {"std": [40], "mean": [1000]}}
lat, lon = torch.meshgrid(torch.linspace(60, 0, 25), torch.linspace(100, 180, 33), indexing="ij")
rlat, rlon = torch.meshgrid(torch.linspace(30, 10, 121), torch.linspace(120, 140, 121), indexing="ij")
static = lambda y, x: torch.stack((y/90, x/180-1, torch.zeros_like(y), torch.zeros_like(y)))[None]
self.inputs = {"global_static": static(lat, lon), "regional_static": static(rlat, rlon),
"center": torch.tensor([[20., 130.]])}
self.outputs = []
for i in range(20):
y, x = torch.meshgrid(torch.linspace(25, 15, 65), torch.linspace(125+i*.5, 135+i*.5, 65), indexing="ij")
core = -torch.exp(-((x-(130+i*.5))**2+(y-20)**2)/3)[None, None]
mask = torch.ones_like(core, dtype=torch.bool)
mask[:, :, :, 0] = False
self.outputs.append({"core": core, "core_lat": y[None], "core_lon": x[None], "core_valid": mask,
"regional_valid": torch.ones((1,1,121,121),dtype=torch.bool), "track_valid": torch.tensor([True]),
"pressure": torch.tensor([800.])})
self.issue = 1536624000 * 10**9
def test_core_units_geography_and_every_valid_time(self):
result = export_pressure_fields(self.outputs, self.inputs, self.contract, self.issue)
self.assertEqual(result["core_mslp_hpa"].shape, (20,65,65))
self.assertEqual(result["core_latitude_deg"].shape, (20,65,65))
self.assertEqual(result["core_longitude_deg"].shape, (20,65,65))
self.assertEqual(result["core_valid"].dtype, np.dtype(bool))
self.assertEqual(result["regional_valid"].shape, (20,121,121))
np.testing.assert_array_equal(result["valid_time_ns"], self.issue+np.arange(6,121,6,dtype=np.int64)*3600*10**9)
self.assertEqual(int(result["member_count"]), 1)
self.assertEqual(float(result["core_information_spacing_km"]), 20)
self.assertAlmostEqual(float(result["core_mslp_hpa"][0,32,32]),960)
np.testing.assert_array_equal(result["core_longitude_deg"][19], self.outputs[19]["core_lon"][0].numpy())
self.assertFalse(result["core_valid"][:,:,0].any())
def test_capture_preserves_tensors_and_never_inserts_scalar_pressure(self):
before = [{k:v.clone() for k,v in o.items()} for o in self.outputs]
result = export_pressure_fields(self.outputs, self.inputs, self.contract, self.issue)
self.assertGreater(float(result["core_mslp_hpa"].min()), float(self.outputs[0]["pressure"][0]))
for old, new in zip(before, self.outputs):
for key in old:
self.assertTrue(torch.equal(old[key], new[key]), key)
def test_a_moving_core_can_leave_the_original_fixed_patch(self):
result = export_pressure_fields(self.outputs, self.inputs, self.contract, self.issue)
fixed_east = float(result["regional_longitude_deg"].max())
self.assertGreater(float(result["core_longitude_deg"][-1].max()), fixed_east)
np.testing.assert_array_equal(result["regional_longitude_deg"], (self.inputs["regional_static"][0,1].numpy()+1)*180)
def test_incomplete_rollout_is_rejected(self):
with self.assertRaisesRegex(ValueError, "twenty"):
export_pressure_fields(self.outputs[:-1], self.inputs, self.contract, self.issue)
def plot_data(self):
result = export_pressure_fields(self.outputs, self.inputs, self.contract, self.issue)
result.update(lead_hours=np.arange(6,121,6), basin_mslp_hpa=np.stack([
1005+self.inputs["global_static"][0,0].numpy()*10+i*.1 for i in range(20)]),
central_pressure_hpa=np.full(20,960),
track_lat_lon=np.column_stack((np.full(20,20),130+np.arange(20)*.5)))
return result
def test_renderer_accepts_real_geography_and_distinct_leads(self):
with tempfile.TemporaryDirectory(prefix="pressure-export-test-") as folder:
first, last = Path(folder)/"first.png", Path(folder)/"last.png"
data = self.plot_data()
render_pressure_map(data, first, 6, 4)
render_pressure_map(data, last, 120, 4)
self.assertGreater(first.stat().st_size, 10000)
self.assertNotEqual(hashlib.sha256(first.read_bytes()).digest(), hashlib.sha256(last.read_bytes()).digest())
def test_missing_core_is_masked_not_filled_from_scalar(self):
data = self.plot_data()
data["core_valid"][:] = False
original = data["core_mslp_hpa"].copy()
with tempfile.TemporaryDirectory(prefix="pressure-export-test-") as folder:
render_pressure_map(data, Path(folder)/"masked.png", 120)
np.testing.assert_array_equal(data["core_mslp_hpa"], original)
def test_invalid_contour_interval_and_unregistered_ensemble_are_rejected(self):
for interval in (0,-1,np.nan):
with self.assertRaisesRegex(ValueError, "interval"):
render_pressure_map(self.plot_data(), "unused.png", interval=interval)
data = self.plot_data()
data["member_count"] = np.array(50)
with self.assertRaisesRegex(ValueError, "one-member"):
render_pressure_map(data, "unused.png")
def test_released_neural_source_hashes_are_unchanged(self):
directory = ROOT/"models/trackformer_1_2_field"
manifest = json.loads((directory/"manifest.json").read_text())
for name, expected in manifest["source_module_sha256"].items():
self.assertEqual(hashlib.sha256((directory/name).read_bytes()).hexdigest(), expected, name)
if __name__ == "__main__":
unittest.main()
|