92 lines
3.6 KiB
Python
92 lines
3.6 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Iterable, Mapping
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SiteState:
|
|
site_id: str
|
|
region: str
|
|
on_hand: float
|
|
forecast_demand: float
|
|
equity_weight: float = 1.0
|
|
continuity_weight: float = 1.0
|
|
triage_priority: float = 1.0
|
|
minimum_service: float = 0.0
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class HumanitarianPolicy:
|
|
equity_weight: float = 1.0
|
|
continuity_weight: float = 1.0
|
|
triage_weight: float = 1.0
|
|
regional_floor: float = 0.0
|
|
|
|
def site_prior(self, site: SiteState) -> float:
|
|
return (
|
|
site.equity_weight * self.equity_weight
|
|
+ site.continuity_weight * self.continuity_weight
|
|
+ site.triage_priority * self.triage_weight
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class AllocationResult:
|
|
allocations: dict[str, float]
|
|
unmet_demand: dict[str, float]
|
|
total_allocated: float
|
|
iterations: int
|
|
|
|
|
|
class DistributedAllocator:
|
|
def __init__(self, iterations: int = 8, step_size: float = 0.5) -> None:
|
|
self.iterations = iterations
|
|
self.step_size = step_size
|
|
|
|
@staticmethod
|
|
def _project_to_sum(values: list[float], total: float, lowers: list[float], uppers: list[float]) -> list[float]:
|
|
if not values:
|
|
return []
|
|
low = min(v - u for v, u in zip(values, uppers)) - abs(total)
|
|
high = max(v - l for v, l in zip(values, lowers)) + abs(total)
|
|
for _ in range(80):
|
|
shift = (low + high) / 2.0
|
|
projected = [min(upper, max(lower, value - shift)) for value, lower, upper in zip(values, lowers, uppers)]
|
|
current = sum(projected)
|
|
if abs(current - total) < 1e-8:
|
|
return projected
|
|
if current > total:
|
|
low = shift
|
|
else:
|
|
high = shift
|
|
shift = (low + high) / 2.0
|
|
return [min(upper, max(lower, value - shift)) for value, lower, upper in zip(values, lowers, uppers)]
|
|
|
|
def allocate(self, sites: Iterable[SiteState], total_available: float, policy: HumanitarianPolicy | None = None) -> AllocationResult:
|
|
policy = policy or HumanitarianPolicy()
|
|
sites = list(sites)
|
|
if not sites:
|
|
return AllocationResult({}, {}, 0.0, 0)
|
|
|
|
lowers = [max(0.0, site.minimum_service) for site in sites]
|
|
uppers = [max(lower, min(site.forecast_demand, site.on_hand + total_available)) for site, lower in zip(sites, lowers)]
|
|
weights = [max(0.1, policy.site_prior(site)) for site in sites]
|
|
total = min(total_available, sum(uppers))
|
|
x = [min(upper, max(lower, total * weight / sum(weights))) for weight, lower, upper in zip(weights, lowers, uppers)]
|
|
duals = [0.0 for _ in sites]
|
|
|
|
for _ in range(self.iterations):
|
|
proposals = []
|
|
for idx, site in enumerate(sites):
|
|
target = min(site.forecast_demand, site.on_hand + total_available)
|
|
prior = policy.site_prior(site)
|
|
local_update = x[idx] + self.step_size * ((target - x[idx]) / max(1.0, target) + prior - duals[idx])
|
|
proposals.append(local_update)
|
|
x = self._project_to_sum(proposals, total, lowers, uppers)
|
|
duals = [dual + proposal - alloc for dual, proposal, alloc in zip(duals, proposals, x)]
|
|
|
|
allocations = {site.site_id: round(value, 4) for site, value in zip(sites, x)}
|
|
unmet = {site.site_id: round(max(0.0, site.forecast_demand - allocations[site.site_id]), 4) for site in sites}
|
|
return AllocationResult(allocations=allocations, unmet_demand=unmet, total_allocated=round(sum(x), 4), iterations=self.iterations)
|