only support l1 and l2 norms, fix hessian

zap/devices/injector.py

@@ -60,14 +60,10 @@ def scale_power(self, scale):
     # CORE MODELING FUNCTIONS
     # ====
 
-    def equality_constraints(
-        self, power, angle, _, nominal_capacity=None, la=np, envelope=None
-    ):
+    def equality_constraints(self, power, angle, _, nominal_capacity=None, la=np, envelope=None):
         return []
 
-    def inequality_constraints(
-        self, power, angle, _, nominal_capacity=None, la=np, envelope=None
-    ):
+    def inequality_constraints(self, power, angle, _, nominal_capacity=None, la=np, envelope=None):
         nominal_capacity = self.parameterize(nominal_capacity=nominal_capacity, la=la)
         power = power[0]
 
@@ -76,9 +72,7 @@ def inequality_constraints(
             power - la.multiply(self.max_power, nominal_capacity),
         ]
 
-    def operation_cost(
-        self, power, angle, _, nominal_capacity=None, la=np, envelope=None
-    ):
+    def operation_cost(self, power, angle, _, nominal_capacity=None, la=np, envelope=None):
         nominal_capacity = self.parameterize(nominal_capacity=nominal_capacity, la=la)
         power = power[0] - la.multiply(self.min_power, nominal_capacity)
 
@@ -118,7 +112,7 @@ def _hessian_power(self, hessians, power, angle, _, nominal_capacity=None, la=np
         if self.quadratic_cost is None:
             return hessians
 
-        hessians[0] += 2 * sp.diags((self.quadratic_cost * power[0]).ravel())
+        hessians[0] += 2 * sp.diags((self.quadratic_cost * np.ones_like(power[0])).ravel())
         return hessians
 
     # ====
@@ -142,9 +136,7 @@ def admm_prox_update(
 
         if self.has_changed:
             quadratic_cost = (
-                0.0 * self.linear_cost
-                if self.quadratic_cost is None
-                else self.quadratic_cost
+                0.0 * self.linear_cost if self.quadratic_cost is None else self.quadratic_cost
             )
             pmax = torch.multiply(self.max_power, nominal_capacity)
             pmin = torch.multiply(self.min_power, nominal_capacity)
@@ -153,9 +145,7 @@ def admm_prox_update(
 
         quadratic_cost, pmax, pmin = self.admm_data
 
-        return _admm_prox_update(
-            power, rho_power, self.linear_cost, quadratic_cost, pmin, pmax
-        )
+        return _admm_prox_update(power, rho_power, self.linear_cost, quadratic_cost, pmin, pmax)
 
     def get_admm_power_weights(self, power, strategy: str, nominal_capacity=None):
         nominal_capacity = self.parameterize(nominal_capacity=nominal_capacity)
@@ -194,12 +184,8 @@ class Generator(AbstractInjector):
     quadratic_cost: Optional[NDArray] = field(default=None, converter=make_dynamic)
     capital_cost: Optional[NDArray] = field(default=None, converter=make_dynamic)
     emission_rates: Optional[NDArray] = field(default=None, converter=make_dynamic)
-    min_nominal_capacity: Optional[NDArray] = field(
-        default=None, converter=make_dynamic
-    )
-    max_nominal_capacity: Optional[NDArray] = field(
-        default=None, converter=make_dynamic
-    )
+    min_nominal_capacity: Optional[NDArray] = field(default=None, converter=make_dynamic)
+    max_nominal_capacity: Optional[NDArray] = field(default=None, converter=make_dynamic)
 
     # TODO - Add dimension checks
 
@@ -280,9 +266,7 @@ def sample_time(self, time_periods, original_time_horizon):
 
 
 @torch.jit.script
-def _admm_prox_update(
-    power: list[torch.Tensor], rho: float, lin_cost, quad_cost, pmin, pmax
-):
+def _admm_prox_update(power: list[torch.Tensor], rho: float, lin_cost, quad_cost, pmin, pmax):
     # Problem is
     #     min_p    a (p - pmin)^2 + b (p - pmin) + (rho / 2) || (p - power) ||_2^2 + {box constraints}
     # Objective derivative is

zap/devices/power_target.py

@@ -2,9 +2,8 @@
 import scipy.sparse as sp
 import cvxpy
 import torch
-from attrs import define, field, Factory
+from attrs import define, field
 
-from typing import Optional
 from numpy.typing import NDArray
 
 from .abstract import AbstractDevice, get_time_horizon, make_dynamic
@@ -19,6 +18,9 @@ class PowerTarget(AbstractDevice):
     target_power: NDArray = field(converter=make_dynamic)
     norm_order: int = field(default=2)
 
+    def __post_init__(self):
+        assert self.norm_order in [1, 2]
+
     @property
     def terminals(self):
         return self.terminal
@@ -40,12 +42,12 @@ def inequality_constraints(self, power, angle, _, target_power=None, la=np):
     def operation_cost(self, power, angle, _, target_power=None, la=np):
         target_power = self.parameterize(target_power=target_power, la=la)
 
-        if la == torch:
-            return 0.5 * torch.linalg.vector_norm(power[0] - target_power, ord=self.norm_order)
-        if la == np:
-            return 0.5 * np.linalg.norm((power[0] - target_power).reshape(-1), ord=self.norm_order)
-        if la == cvxpy:
-            return 0.5 * cvxpy.norm(power[0] - target_power, self.norm_order)
+        err = power[0] - target_power
+        if self.norm_order == 1:
+            return la.sum(la.abs(err))
+
+        else:  # L2
+            return (0.5) * la.sum(la.square(err))
 
     # ====
     # DIFFERENTIATION
@@ -60,7 +62,11 @@ def _inequality_matrices(self, inequalities, target_power=None, la=np):
     def _hessian_power(self, hessians, power, angle, _, target_power=None, la=np):
         target_power = self.parameterize(target_power=target_power, la=la)
 
-        hessians[0] += sp.diags((power[0] - target_power).ravel())
+        if self.norm_order == 2:
+            hessians[0] += sp.diags(np.ones_like(power[0]).ravel())
+        else:  # L1
+            pass
+
         return hessians
 
     # ====