[Feat] Add YZLSM Solver (#10)

endure/lsm/solver/__init__.py

@@ -1,3 +1,4 @@
 from .classic_solver import ClassicSolver
 from .qlsm_solver import QLSMSolver
 from .klsm_solver import KLSMSolver
+from .yzlsm_solver import YZLSMSolver

endure/lsm/solver/yzlsm_solver.py

@@ -0,0 +1,112 @@
+from typing import Any, Optional, Callable, Tuple
+
+import numpy as np
+import scipy.optimize as SciOpt
+
+from endure.lsm.cost import EndureCost
+from endure.lsm.types import LSMDesign, Policy, System
+from .util import kl_div_con, get_bounds
+from .util import (
+    H_DEFAULT,
+    T_DEFAULT,
+    LAMBDA_DEFAULT,
+    ETA_DEFAULT,
+    Y_DEFAULT,
+    Z_DEFAULT,
+)
+
+
+class YZLSMSolver:
+    def __init__(self, config: dict[str, Any]):
+        self.config = config
+        self.cf = EndureCost(config["lsm"]["max_levels"])
+
+    def robust_objective(
+        self,
+        x: np.ndarray,
+        system: System,
+        rho: float,
+        z0: float,
+        z1: float,
+        q: float,
+        w: float,
+    ) -> float:
+        h, t, y, z, lamb, eta = x
+        design = LSMDesign(h=h, T=t, Y=y, Z=z, policy=Policy.YZHybrid)
+        query_cost = 0
+        query_cost += z0 * kl_div_con((self.cf.Z0(design, system) - eta) / lamb)
+        query_cost += z1 * kl_div_con((self.cf.Z1(design, system) - eta) / lamb)
+        query_cost += q * kl_div_con((self.cf.Q(design, system) - eta) / lamb)
+        query_cost += w * kl_div_con((self.cf.W(design, system) - eta) / lamb)
+        cost = eta + (rho * lamb) + (lamb * query_cost)
+
+        return cost
+
+    def nominal_objective(
+        self,
+        x: np.ndarray,
+        system: System,
+        z0: float,
+        z1: float,
+        q: float,
+        w: float,
+    ) -> float:
+        h, t, y, z = x
+        design = LSMDesign(h=h, T=t, Y=y, Z=z, policy=Policy.YZHybrid)
+        cost = self.cf.calc_cost(design, system, z0, z1, q, w)
+
+        return cost
+
+    def get_robust_design(
+        self,
+        system: System,
+        rho: float,
+        z0: float,
+        z1: float,
+        q: float,
+        w: float,
+        init_args: np.ndarray = np.array(
+            [H_DEFAULT, T_DEFAULT, LAMBDA_DEFAULT, ETA_DEFAULT]
+        ),
+        minimizer_kwargs: dict = {},
+        callback_fn: Optional[Callable] = None,
+    ) -> Tuple[LSMDesign, SciOpt.OptimizeResult]:
+        raise NotImplementedError
+
+    def get_nominal_design(
+        self,
+        system: System,
+        z0: float,
+        z1: float,
+        q: float,
+        w: float,
+        init_args: np.ndarray = np.array([H_DEFAULT, T_DEFAULT, Y_DEFAULT, Z_DEFAULT]),
+        minimizer_kwargs: dict = {},
+        callback_fn: Optional[Callable] = None,
+    ) -> Tuple[LSMDesign, SciOpt.OptimizeResult]:
+        default_kwargs = {
+            "method": "SLSQP",
+            "bounds": get_bounds(
+                config=self.config,
+                system=system,
+                robust=False,
+            ),
+            "options": {"ftol": 1e-6, "disp": False, "maxiter": 1000},
+        }
+        default_kwargs.update(minimizer_kwargs)
+
+        solution = SciOpt.minimize(
+            fun=lambda x: self.nominal_objective(x, system, z0, z1, q, w),
+            x0=init_args,
+            callback=callback_fn,
+            **default_kwargs
+        )
+        design = LSMDesign(
+            h=solution.x[0],
+            T=solution.x[1],
+            Y=solution.x[2],
+            Z=solution.x[3],
+            policy=Policy.YZHybrid,
+        )
+
+        return design, solution

endure/lsm/types.py

@@ -1,5 +1,6 @@
 from dataclasses import dataclass, field
 import enum
+from typing import Tuple
 
 
 class Policy(enum.Enum):
@@ -42,6 +43,13 @@ class LSMDesign:
     K: list[float] = field(default_factory=list)
 
 
+@dataclass
+class LSMBounds:
+    max_levels: int = 20
+    bpe: Tuple[float, float] = (1.0, 10.0)
+    size_ratio: Tuple[float, float] = (2, 31)
+
+
 @dataclass
 class Workload:
     z0: float = 0.25

endure/ltune/util/ltune_eval.py

@@ -96,6 +96,8 @@ def get_solver_nominal_design(
             solver = Solver.QLSMSolver(self.config)
         elif self.design_type == "KLSM":
             solver = Solver.KLSMSolver(self.config)
+        elif self.design_type == "YZLSM":
+            solver = Solver.YZLSMSolver(self.config)
         else: # design_type == "Classic"
             solver = Solver.ClassicSolver(self.config)
 

jobs/create_tunings.py

@@ -27,6 +27,7 @@ def _generate_solver(self):
         choices = {
             "KCost": Solvers.KLSMSolver(self.config),
             "QCost": Solvers.QLSMSolver(self.config),
+            "YZCost": Solvers.YZLSMSolver(self.config),
             "TierCost": Solvers.ClassicSolver(self.config, policies=[Policy.Tiering]),
             "LevelCost": Solvers.ClassicSolver(self.config, policies=[Policy.Leveling]),
             "TierLevelCost": Solvers.ClassicSolver(self.config)