| import random |
| import uuid |
| from typing import List, Optional |
| import sys |
| import os |
| sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) |
|
|
| from models import ( |
| DistrictObservation, |
| DistrictTruth, |
| CityObservation, |
| CityState, |
| ) |
| from server.constants import ( |
| SPREAD_RATE_MIN, |
| SPREAD_RATE_MAX, |
| GROWTH_HINT_NOISE, |
| INFECTION_THRESHOLD, |
| SAFE_THRESHOLD, |
| HOSPITAL_BREACH_POINT, |
| NATURAL_RECOVERY_RATE, |
| ) |
|
|
|
|
| def generate_districts( |
| num_districts: int, |
| seed_infections: List[float], |
| ) -> List[DistrictTruth]: |
| assert len(seed_infections) == num_districts |
| raw_densities = [random.uniform(0.5, 1.5) for _ in range(num_districts)] |
| total = sum(raw_densities) |
| densities = [round(d / total, 4) for d in raw_densities] |
| districts = [] |
| for i in range(num_districts): |
| districts.append(DistrictTruth( |
| district_id = i, |
| true_infection_rate = seed_infections[i], |
| true_spread_rate = round(random.uniform(SPREAD_RATE_MIN, SPREAD_RATE_MAX), 4), |
| hospital_capacity_remaining = 1.0, |
| population_density = densities[i], |
| days_since_tested = 99, |
| restriction_active = False, |
| deployed_resources = 0, |
| )) |
| return districts |
|
|
|
|
| def generate_episode_id() -> str: |
| return str(uuid.uuid4())[:8] |
|
|
|
|
| def build_observation( |
| state: CityState, |
| step_count: int, |
| reward: Optional[float] = None, |
| message: Optional[str] = None, |
| done: bool = False, |
| ) -> CityObservation: |
| """ |
| Build the agent-visible observation from hidden city state. |
| Hard task enforces a 3-day lag on reported infection rates. |
| Hospital capacity and growth hints are always real-time. |
| """ |
| district_observations = [] |
| for i, district in enumerate(state.districts): |
| if state.data_lag_days > 0 and len(state.infection_history) >= state.data_lag_days: |
| reported_rate = state.infection_history[-state.data_lag_days][i] |
| else: |
| reported_rate = district.true_infection_rate |
|
|
| noise = random.uniform(-GROWTH_HINT_NOISE, GROWTH_HINT_NOISE) |
| growth_hint = round(max(0.0, min(1.0, district.true_spread_rate + noise)), 4) |
|
|
| district_observations.append(DistrictObservation( |
| district_id = district.district_id, |
| reported_infection_rate = round(reported_rate, 4), |
| growth_rate_hint = growth_hint, |
| hospital_capacity_remaining = round(district.hospital_capacity_remaining, 4), |
| population_density = district.population_density, |
| tested_recently = district.days_since_tested <= 2, |
| restriction_active = district.restriction_active, |
| )) |
|
|
| return CityObservation( |
| districts = district_observations, |
| available_resources = state.available_resources, |
| current_step = step_count, |
| max_steps = state.max_steps, |
| data_lag_days = state.data_lag_days, |
| done = done, |
| reward = reward, |
| message = message, |
| ) |
|
|
|
|
| def compute_spread(districts: List[DistrictTruth]) -> List[float]: |
| """ |
| Advance infection rates by one day using a simplified SIR-inspired model. |
| |
| Net change per district: |
| delta = effective_spread_rate - natural_recovery + geographic_spillover |
| |
| Spillover is linear (no wrap-around). District 0 and the last district |
| are not adjacent, which mirrors a city corridor layout rather than a ring. |
| """ |
| from server.constants import ( |
| ALLOCATE_REDUCTION, |
| RESTRICT_REDUCTION, |
| SPILLOVER_RATE, |
| NATURAL_RECOVERY_RATE, |
| ) |
|
|
| n = len(districts) |
| new_rates = [] |
|
|
| for i, district in enumerate(districts): |
| effective_spread = district.true_spread_rate |
|
|
| if district.restriction_active: |
| effective_spread = max(0.0, effective_spread - RESTRICT_REDUCTION) |
|
|
| if district.deployed_resources > 0: |
| effective_spread = max( |
| 0.0, |
| effective_spread - (ALLOCATE_REDUCTION * district.deployed_resources) |
| ) |
|
|
| net_change = effective_spread - NATURAL_RECOVERY_RATE |
| new_rate = district.true_infection_rate + net_change |
|
|
| if i > 0: |
| new_rate += districts[i - 1].true_infection_rate * SPILLOVER_RATE |
| if i < n - 1: |
| new_rate += districts[i + 1].true_infection_rate * SPILLOVER_RATE |
|
|
| new_rates.append(round(min(1.0, max(0.0, new_rate)), 4)) |
|
|
| return new_rates |
|
|
|
|
| def get_highest_infected_district(districts: List[DistrictTruth]) -> int: |
| return max(districts, key=lambda d: d.true_infection_rate).district_id |
|
|
|
|
| def all_districts_contained(districts: List[DistrictTruth]) -> bool: |
| return all(d.true_infection_rate < SAFE_THRESHOLD for d in districts) |
|
|
|
|
| def any_hospital_breached(districts: List[DistrictTruth]) -> bool: |
| return any(d.hospital_capacity_remaining <= HOSPITAL_BREACH_POINT for d in districts) |
|
|
|
|
| def districts_above_threshold(districts: List[DistrictTruth]) -> List[DistrictTruth]: |
| return [d for d in districts if d.true_infection_rate > INFECTION_THRESHOLD] |
|
|
|
|
| def snapshot_infection_rates(districts: List[DistrictTruth]) -> List[float]: |
| return [d.true_infection_rate for d in sorted(districts, key=lambda d: d.district_id)] |
|
|
