83 lines
3.1 KiB
Python
83 lines
3.1 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
import math
|
|
|
|
import numpy as np
|
|
|
|
from .models import LocalProblem, PlanAction, PlanDelta, SolverResult
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _State:
|
|
x: np.ndarray
|
|
z: np.ndarray
|
|
u: np.ndarray
|
|
|
|
|
|
class ADMMLiteSolver:
|
|
def __init__(self, *, rho: float = 1.0, tolerance: float = 1e-6, max_iterations: int = 100) -> None:
|
|
if rho <= 0:
|
|
raise ValueError("rho must be positive")
|
|
self.rho = rho
|
|
self.tolerance = tolerance
|
|
self.max_iterations = max_iterations
|
|
|
|
def solve(self, problems: list[LocalProblem]) -> SolverResult:
|
|
if not problems:
|
|
raise ValueError("at least one local problem is required")
|
|
|
|
horizons = {len(problem.preferred_dispatch) for problem in problems}
|
|
if len(horizons) != 1:
|
|
raise ValueError("all problems must share the same dispatch horizon")
|
|
|
|
horizon = horizons.pop()
|
|
lower = np.array([problem.lower_bounds or [-math.inf] * horizon for problem in problems], dtype=float)
|
|
upper = np.array([problem.upper_bounds or [math.inf] * horizon for problem in problems], dtype=float)
|
|
preferred = np.array([problem.preferred_dispatch for problem in problems], dtype=float)
|
|
weights = np.array([problem.quadratic_weight for problem in problems], dtype=float).reshape(-1, 1)
|
|
|
|
x = preferred.copy()
|
|
z = preferred.mean(axis=0)
|
|
u = np.zeros_like(x)
|
|
primal_residual = float("inf")
|
|
dual_residual = float("inf")
|
|
|
|
for iteration in range(1, self.max_iterations + 1):
|
|
previous_z = z.copy()
|
|
x = np.clip((weights * preferred + self.rho * (z - u)) / (weights + self.rho), lower, upper)
|
|
z = np.mean(x + u, axis=0)
|
|
u = u + x - z
|
|
|
|
primal_residual = float(np.max(np.linalg.norm(x - z, axis=1)))
|
|
dual_residual = float(np.linalg.norm(z - previous_z) * math.sqrt(len(problems)))
|
|
if primal_residual <= self.tolerance and dual_residual <= self.tolerance:
|
|
break
|
|
|
|
deltas: list[PlanDelta] = []
|
|
for index, problem in enumerate(problems):
|
|
actions = [
|
|
PlanAction(resource_id=f"{problem.site_id}:{slot}", action="dispatch", value=float(value), metadata={"slot": slot})
|
|
for slot, value in enumerate(x[index])
|
|
]
|
|
deltas.append(
|
|
PlanDelta(
|
|
site_id=problem.site_id,
|
|
adapter_id=f"adapter:{problem.site_id}",
|
|
revision=iteration,
|
|
round_id=problem.round_id,
|
|
parent_revision=max(0, iteration - 1),
|
|
actions=actions,
|
|
metadata={"solver": "admm-lite", "round_id": problem.round_id},
|
|
)
|
|
)
|
|
|
|
return SolverResult(
|
|
consensus=[float(value) for value in z],
|
|
iterations=iteration,
|
|
primal_residual=primal_residual,
|
|
dual_residual=dual_residual,
|
|
deltas=deltas,
|
|
metadata={"rho": self.rho, "tolerance": self.tolerance, "max_iterations": self.max_iterations},
|
|
)
|