idea150-medimesh-federated-.../src/medimesh/solver.py

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)