test_pyropredict.py

"""Unit tests for PyroPredict."""

import numpy as np
import pytest

from pyropredict import (
    synthesize_california_chaparral, synthesize_boreal_forest,
    generate_santa_ana, generate_typical_summer,
    CellularAutomatonSpread, rothermel_rate_of_spread,
    IGNITED, BURNT, UNBURNT, UNBURNABLE,
    HybridForecaster,
    RoadNetwork, evacuation_route, time_evacuated_fraction,
    optimal_crew_placement,
    FUEL_MODELS,
)


def test_rothermel_zero_when_no_fuel():
    fm = FUEL_MODELS["nonfuel"]
    r = rothermel_rate_of_spread(
        fm=fm, moisture_pct=10, wind_speed_kmh=30, wind_from_deg=0,
        direction_deg=180, slope_deg=5, slope_aspect_deg=180,
    )
    assert r == 0


def test_rothermel_zero_when_moisture_above_ext():
    fm = FUEL_MODELS["grass"]
    r = rothermel_rate_of_spread(
        fm=fm, moisture_pct=fm.moisture_ext + 1,
        wind_speed_kmh=30, wind_from_deg=0,
        direction_deg=180, slope_deg=0, slope_aspect_deg=0,
    )
    assert r == 0


def test_rothermel_wind_amplifies_downwind():
    """ROS in the wind-aligned direction should exceed ROS against the wind."""
    fm = FUEL_MODELS["chaparral"]
    # Wind from north (0), toward south (180) -- propagation in 180 should be fastest
    r_down = rothermel_rate_of_spread(
        fm=fm, moisture_pct=10, wind_speed_kmh=50, wind_from_deg=0,
        direction_deg=180, slope_deg=0, slope_aspect_deg=0,
    )
    r_up = rothermel_rate_of_spread(
        fm=fm, moisture_pct=10, wind_speed_kmh=50, wind_from_deg=0,
        direction_deg=0, slope_deg=0, slope_aspect_deg=0,
    )
    assert r_down > 3 * r_up


def test_rothermel_slope_amplifies_uphill():
    fm = FUEL_MODELS["chaparral"]
    r_up = rothermel_rate_of_spread(
        fm=fm, moisture_pct=10, wind_speed_kmh=5, wind_from_deg=0,
        direction_deg=180, slope_deg=30, slope_aspect_deg=180,
    )
    r_down = rothermel_rate_of_spread(
        fm=fm, moisture_pct=10, wind_speed_kmh=5, wind_from_deg=0,
        direction_deg=0, slope_deg=30, slope_aspect_deg=180,
    )
    assert r_up > 1.5 * r_down


def test_ca_simulate_spread_with_wind():
    """Under strong wind the fire should spread further than without wind."""
    land = synthesize_california_chaparral(seed=0)
    ca = CellularAutomatonSpread(land)
    w_calm = generate_typical_summer(n_hours=8, mean_wind_kmh=5, seed=0)
    w_high = generate_santa_ana(n_hours=8, mean_wind_kmh=60, seed=0)
    res_calm = ca.simulate([(35, 35)], w_calm, n_hours=4, dt_min=20, seed=0)
    res_high = ca.simulate([(35, 35)], w_high, n_hours=4, dt_min=20, seed=0)
    burnt_calm = int(((res_calm.state == IGNITED) | (res_calm.state == BURNT)).sum())
    burnt_high = int(((res_high.state == IGNITED) | (res_high.state == BURNT)).sum())
    assert burnt_high > burnt_calm


def test_ca_unburnable_cells_stay_unburnable():
    land = synthesize_california_chaparral(seed=0)
    ca = CellularAutomatonSpread(land)
    w = generate_santa_ana(n_hours=4, seed=0)
    res = ca.simulate([(35, 35)], w, n_hours=2, dt_min=20, seed=0)
    # All water cells must remain UNBURNABLE
    H, W = land.shape()
    for r in range(H):
        for c in range(W):
            if land.is_water[r, c]:
                assert res.state[r, c] == UNBURNABLE


def test_hybrid_forecaster_increases_burned_area():
    """When the CA underestimates burned area (typical bias), the hybrid
    forecaster should produce a larger forecast that's closer to truth."""
    land = synthesize_california_chaparral(seed=0)
    w = generate_santa_ana(n_hours=8, seed=0)
    ca = CellularAutomatonSpread(land)
    # Simulate "ground truth" with longer residence time -> more spread
    truth = ca.simulate([(35, 35)], w, n_hours=8, dt_min=15, seed=0)
    # Pretend the truth has additional spotting / underestimation:
    # we manually expand the burnt area by neighbouring cells
    truth_state = truth.state.copy()
    H, W = land.shape()
    for _ in range(2):
        burning = (truth_state == IGNITED) | (truth_state == BURNT)
        # 4-neighbour dilation
        ext = (np.roll(burning, 1, 0) | np.roll(burning, -1, 0) |
               np.roll(burning, 1, 1) | np.roll(burning, -1, 1))
        new = ext & (truth_state == UNBURNT)
        truth_state[new] = BURNT
    from pyropredict.spread import SpreadState
    truth_aug = SpreadState(state=truth_state, ignition_time=truth.ignition_time,
                              intensity=truth.intensity)
    hyb = HybridForecaster(land).fit(
        ignition_cells=[(35, 35)], weather=w,
        true_state=truth_aug, n_hours=8, dt_min=15, seed=0,
    )
    pred = hyb.forecast([(35, 35)], w, n_hours=8, dt_min=15, seed=42)
    ca_pred = ca.simulate([(35, 35)], w, n_hours=8, dt_min=15, seed=42)
    ca_burned = int(((ca_pred.state == IGNITED) | (ca_pred.state == BURNT)).sum())
    hyb_burned = int(((pred.state == IGNITED) | (pred.state == BURNT)).sum())
    truth_burned = int(((truth_aug.state == IGNITED) | (truth_aug.state == BURNT)).sum())
    # Hybrid should be between CA and truth
    assert hyb_burned >= ca_burned * 0.8           # hybrid never dramatically worse
    # If multiplier > 1, hybrid burns MORE (closer to truth)
    if hyb.ignition_prob_multiplier > 1.05:
        assert hyb_burned > ca_burned


def test_evacuation_route_exists_for_road_source():
    land = synthesize_california_chaparral(seed=0)
    net = RoadNetwork.from_landscape(land)
    ca = CellularAutomatonSpread(land)
    w = generate_typical_summer(n_hours=4, seed=0)
    res = ca.simulate([(80, 80)], w, n_hours=2, dt_min=20, seed=0)
    # Pick a road cell not yet burning
    road_cell = next(iter(net.nodes))
    r = evacuation_route(net, road_cell, res)
    # Either a route exists or all sinks blocked
    assert isinstance(r, dict)


def test_forecast_aware_routing_avoids_burning_path():
    land = synthesize_california_chaparral(seed=0)
    net = RoadNetwork.from_landscape(land)
    ca = CellularAutomatonSpread(land)
    w = generate_santa_ana(n_hours=4, seed=0)
    truth = ca.simulate([(20, 50)], w, n_hours=2, dt_min=20, seed=0)
    forecast = ca.simulate([(20, 50)], w, n_hours=4, dt_min=20, seed=1)
    # Pick a populated cell as evacuation source
    pops = np.argwhere(land.population > 0)
    if pops.size == 0:
        pytest.skip("no populated cells")
    rt_naive = time_evacuated_fraction(
        net, [(int(r), int(c)) for r, c in pops[:50]],
        [int(land.population[r, c]) for r, c in pops[:50]],
        truth, forecast_aware=False, deadline_min=60,
    )
    rt_aware = time_evacuated_fraction(
        net, [(int(r), int(c)) for r, c in pops[:50]],
        [int(land.population[r, c]) for r, c in pops[:50]],
        truth, spread_forecast=forecast,
        forecast_aware=True, deadline_min=60,
    )
    # Forecast-aware should equal or exceed naive
    assert rt_aware["evacuated"] >= rt_naive["evacuated"] - 5


def test_resource_placement_reduces_loss():
    land = synthesize_california_chaparral(seed=0)
    w = generate_santa_ana(n_hours=4, seed=0)
    ca = CellularAutomatonSpread(land)
    forecast = ca.simulate([(30, 30)], w, n_hours=4, dt_min=20, seed=0)
    result = optimal_crew_placement(land, forecast, n_crews=3)
    assert result["final_loss"] <= result["initial_loss"]


if __name__ == "__main__":
    raise SystemExit(pytest.main([__file__, "-v"]))