from __future__ import annotations import csv import math import sys import unittest from pathlib import Path from tempfile import TemporaryDirectory import geopandas as gpd import numpy as np import rasterio from rasterio.transform import from_origin from shapely.geometry import MultiPolygon, box SCRIPTS_DIR = Path(__file__).resolve().parents[1] / "scripts" sys.path.insert(0, str(SCRIPTS_DIR)) from apply_koppen_metric_to_climate_data import apply_koppen_metric # noqa: E402 from build_county_koppen_metric import ( # noqa: E402 MIXED_CLASS, build_koppen_records, classify, rank_class_shares, ) from common.county_zonal_stats import area_weighted_class_weights # noqa: E402 from common.koppen_legend import DEFAULT_KOPPEN_CODE_MAP # noqa: E402 CFA, CFB, DFB, ET = 14, 15, 26, 29 def write_global_raster(path: Path, cells: dict[tuple[int, int], int]) -> None: """Write a 1-degree global raster; row r covers latitudes 89 - r to 90 - r.""" data = np.zeros((180, 360), dtype=np.uint8) for (row, col), code in cells.items(): data[row, col] = code with rasterio.open( path, "w", driver="GTiff", height=180, width=360, count=1, dtype="uint8", crs="EPSG:4326", transform=from_origin(-180.0, 90.0, 1.0, 1.0), nodata=0, ) as destination: destination.write(data, 1) def write_csv(path: Path, fieldnames: list[str], rows: list[dict[str, str]]) -> None: with path.open("w", encoding="utf-8", newline="") as handle: writer = csv.DictWriter(handle, fieldnames=fieldnames) writer.writeheader() writer.writerows(rows) class ClassifyTests(unittest.TestCase): def test_clear_majority_is_predominant(self) -> None: self.assertEqual(classify([("Dfb", 0.60), ("Dfa", 0.30)]), "Dfb") def test_exact_cutoffs_are_predominant(self) -> None: self.assertEqual(classify([("Csa", 0.50), ("BSk", 0.45)]), "Csa") def test_majority_with_close_runner_up_is_mixed(self) -> None: self.assertEqual(classify([("Dfb", 0.501), ("Dfa", 0.499)]), MIXED_CLASS) def test_no_majority_is_mixed(self) -> None: self.assertEqual(classify([("BSh", 0.29), ("Csa", 0.24), ("Dsb", 0.23)]), MIXED_CLASS) def test_single_class_is_predominant(self) -> None: self.assertEqual(classify([("ET", 1.0)]), "ET") def test_no_valid_cells_is_blank(self) -> None: self.assertEqual(classify([]), "") class RankClassSharesTests(unittest.TestCase): def test_shares_are_ranked_and_ties_go_to_smaller_code(self) -> None: ranked = rank_class_shares({DFB: 2.0, CFA: 2.0, ET: 1.0}, DEFAULT_KOPPEN_CODE_MAP) self.assertEqual([name for name, _ in ranked], ["Cfa", "Dfb", "ET"]) self.assertAlmostEqual(ranked[0][1], 0.4) self.assertAlmostEqual(sum(share for _, share in ranked), 1.0) def test_unknown_code_raises(self) -> None: with self.assertRaises(ValueError): rank_class_shares({99: 1.0}, DEFAULT_KOPPEN_CODE_MAP) def test_no_weights_give_no_classes(self) -> None: self.assertEqual(rank_class_shares({}, DEFAULT_KOPPEN_CODE_MAP), []) class KoppenRasterTestCase(unittest.TestCase): def setUp(self) -> None: self._temp_dir = TemporaryDirectory() self.raster_path = Path(self._temp_dir.name) / "koppen.tif" write_global_raster( self.raster_path, { (59, 80): CFA, # 30..31 N, 100..99 W (59, 81): CFB, # 30..31 N, 99..98 W (49, 80): DFB, # 40..41 N, 100..99 W; 99..98 W is ocean (29, 80): DFB, # 60..61 N (79, 80): CFA, # 10..11 N (38, 358): ET, # 51..52 N, 178..179 E (38, 359): ET, # 51..52 N, 179..180 E (38, 0): CFB, # 51..52 N, 180..179 W }, ) def tearDown(self) -> None: self._temp_dir.cleanup() def weights(self, geometry) -> dict[int, float]: with rasterio.open(self.raster_path) as source: return area_weighted_class_weights(source, geometry) class AreaWeightedClassWeightsTests(KoppenRasterTestCase): def test_partial_cell_counts_by_fraction_inside(self) -> None: weights = self.weights(box(-100.0, 30.0, -98.5, 31.0)) self.assertAlmostEqual(weights[CFB] / weights[CFA], 0.5) def test_cells_are_weighted_by_latitude(self) -> None: weights = self.weights(MultiPolygon([box(-100.0, 60.0, -99.0, 61.0), box(-100.0, 10.0, -99.0, 11.0)])) expected = math.cos(math.radians(60.5)) / math.cos(math.radians(10.5)) self.assertAlmostEqual(weights[DFB] / weights[CFA], expected) def test_ocean_cells_are_excluded(self) -> None: weights = self.weights(box(-100.0, 40.0, -98.0, 41.0)) self.assertEqual(list(weights), [DFB]) def test_islands_on_both_sides_of_the_date_line_are_combined(self) -> None: weights = self.weights(MultiPolygon([box(178.0, 51.0, 180.0, 52.0), box(-180.0, 51.0, -179.0, 52.0)])) self.assertAlmostEqual(weights[ET] / weights[CFB], 2.0) def test_no_valid_cells_return_no_weights(self) -> None: self.assertEqual(self.weights(box(-50.0, 30.0, -49.0, 31.0)), {}) class BuildKoppenRecordsTests(KoppenRasterTestCase): def test_counties_are_predominant_mixed_or_blank(self) -> None: counties = gpd.GeoDataFrame( {"county_fips": ["00001", "00002", "00003"], "county_name": ["A", "B", "C"], "state": ["AA"] * 3}, geometry=[box(-100.0, 40.0, -99.0, 41.0), box(-100.0, 30.0, -98.0, 31.0), box(-50.0, 30.0, -49.0, 31.0)], crs="EPSG:4326", ) records = build_koppen_records(counties, self.raster_path, DEFAULT_KOPPEN_CODE_MAP) self.assertEqual([record["koppenZone"] for record in records], ["Dfb", MIXED_CLASS, ""]) self.assertEqual((records[0]["koppenTopShare"], records[0]["koppenSecondClass"]), ("1.0000", "")) self.assertEqual( (records[1]["koppenTopClass"], records[1]["koppenTopShare"], records[1]["koppenSecondShare"]), ("Cfa", "0.5000", "0.5000"), ) self.assertEqual(records[2]["koppenTopShare"], "") METRIC_TEST_FIELDS = ["countyFips", "koppenZone", "koppenTopClass", "koppenTopShare", "koppenSecondClass"] class ApplyKoppenMetricTests(unittest.TestCase): def setUp(self) -> None: self._temp_dir = TemporaryDirectory() base = Path(self._temp_dir.name) self.climate_data = base / "climate-data.csv" self.metric = base / "koppen.csv" self.climate_fields = ["countyFips", "countyName", "state", "koppenZone", "avgTempF", "source"] write_csv( self.climate_data, self.climate_fields, [ {"countyFips": "01001", "countyName": "A", "state": "AL", "koppenZone": "Cfa", "avgTempF": "64.5", "source": "s"}, {"countyFips": "01003", "countyName": "B", "state": "AL", "koppenZone": "Dfb", "avgTempF": "50.1", "source": "s"}, ], ) self.write_metric( [ {"countyFips": "01001", "koppenZone": "Mixed", "koppenTopClass": "Csb", "koppenTopShare": "0.4800", "koppenSecondClass": "Dsb"}, {"countyFips": "01003", "koppenZone": "Dfb", "koppenTopClass": "Dfb", "koppenTopShare": "0.9000", "koppenSecondClass": "Dfc"}, ] ) def tearDown(self) -> None: self._temp_dir.cleanup() def write_metric(self, rows: list[dict[str, str]]) -> None: write_csv(self.metric, METRIC_TEST_FIELDS, rows) def read_climate(self) -> tuple[list[str], list[dict[str, str]]]: with self.climate_data.open("r", encoding="utf-8", newline="") as handle: reader = csv.DictReader(handle) return list(reader.fieldnames or []), list(reader) def test_stripe_columns_are_added_after_koppen_zone(self) -> None: changes, added_columns = apply_koppen_metric(self.climate_data, self.metric, self.climate_data) fields, _ = self.read_climate() self.assertEqual(added_columns, ["koppenPrimaryClass", "koppenSecondaryClass"]) self.assertEqual( fields, ["countyFips", "countyName", "state", "koppenZone", "koppenPrimaryClass", "koppenSecondaryClass", "avgTempF", "source"], ) self.assertEqual( changes, [ ("01001", "koppenZone", "Cfa", "Mixed"), ("01001", "koppenPrimaryClass", "", "Csb"), ("01001", "koppenSecondaryClass", "", "Dsb"), ], ) def test_only_mixed_counties_get_stripe_classes(self) -> None: apply_koppen_metric(self.climate_data, self.metric, self.climate_data) _, rows = self.read_climate() self.assertEqual( [(row["koppenZone"], row["koppenPrimaryClass"], row["koppenSecondaryClass"]) for row in rows], [("Mixed", "Csb", "Dsb"), ("Dfb", "", "")], ) self.assertEqual([row["avgTempF"] for row in rows], ["64.5", "50.1"]) def test_existing_stripe_columns_are_updated_in_place(self) -> None: apply_koppen_metric(self.climate_data, self.metric, self.climate_data) self.write_metric( [ {"countyFips": "01001", "koppenZone": "Cfa", "koppenTopClass": "Cfa", "koppenTopShare": "0.7000", "koppenSecondClass": "Dfb"}, {"countyFips": "01003", "koppenZone": "Dfb", "koppenTopClass": "Dfb", "koppenTopShare": "0.9000", "koppenSecondClass": "Dfc"}, ] ) changes, added_columns = apply_koppen_metric(self.climate_data, self.metric, self.climate_data) fields, rows = self.read_climate() self.assertEqual(added_columns, []) self.assertEqual(fields.count("koppenPrimaryClass"), 1) self.assertEqual((rows[0]["koppenZone"], rows[0]["koppenPrimaryClass"], rows[0]["koppenSecondaryClass"]), ("Cfa", "", "")) self.assertEqual(len(changes), 3) def test_dry_run_writes_nothing(self) -> None: before = self.climate_data.read_bytes() changes, added_columns = apply_koppen_metric(self.climate_data, self.metric, self.climate_data, dry_run=True) self.assertEqual(len(changes), 3) self.assertEqual(len(added_columns), 2) self.assertEqual(self.climate_data.read_bytes(), before) def test_missing_county_raises_without_writing(self) -> None: self.write_metric( [{"countyFips": "01001", "koppenZone": "Mixed", "koppenTopClass": "Csb", "koppenTopShare": "0.4800", "koppenSecondClass": "Dsb"}] ) before = self.climate_data.read_bytes() with self.assertRaises(ValueError): apply_koppen_metric(self.climate_data, self.metric, self.climate_data) self.assertEqual(self.climate_data.read_bytes(), before) def test_mixed_county_without_classes_raises_without_writing(self) -> None: self.write_metric( [ {"countyFips": "01001", "koppenZone": "Mixed", "koppenTopClass": "", "koppenTopShare": "", "koppenSecondClass": ""}, {"countyFips": "01003", "koppenZone": "Dfb", "koppenTopClass": "Dfb", "koppenTopShare": "0.9000", "koppenSecondClass": "Dfc"}, ] ) before = self.climate_data.read_bytes() with self.assertRaises(ValueError): apply_koppen_metric(self.climate_data, self.metric, self.climate_data) self.assertEqual(self.climate_data.read_bytes(), before) if __name__ == "__main__": unittest.main()