fixing all not-implemented hessian issues (#45)

* fixing all not-implemented hessian issues

* making some tests pass with hessian_approx=exact

* delete the remove laters

---------

Co-authored-by: William Zijie Zhang <william@gridmatic.com>

Author: Transurgeon

cvxpy/atoms/affine/binary_operators.py

@@ -32,7 +32,6 @@
 from cvxpy.atoms.affine.sum import sum as cvxpy_sum
 from cvxpy.constraints.constraint import Constraint
 from cvxpy.error import DCPError
-from cvxpy.expressions.constants.constant import Constant
 from cvxpy.expressions.constants.parameter import (
     is_param_affine,
     is_param_free,
@@ -235,33 +234,60 @@ def _grad(self, values):
 
         return [DX, DY]
 
-    def _hess(self, values):
-        """Compute the Hessian of elementwise multiplication w.r.t. each argument.
+    def _verify_hess_vec_args(self):
+        x = self.args[0]
+        y = self.args[1]
+        if x.size != y.size:
+            return False
 
-        For z = x * y (elementwise), returns:
-        - d2z/dx2 = diag(y)
-        - d2z/dy2 = diag(x)
+        if x.is_constant() and y.is_constant():
+            return False
 
-        Args:
-            values: A list of numeric values for the arguments [x, y].
+        # one of the following must be true:
+        # 1. both arguments are variables
+        # 2. one argument is a constant
+        # 3. one argument is a Promote of a variable and the other is a variable
+        both_are_variables = isinstance(x, Variable) and isinstance(y, Variable)
+        one_is_constant = x.is_constant() or y.is_constant()
+        x_is_promote = type(x) == Promote and isinstance(y, Variable)
+        y_is_promote = type(y) == Promote and isinstance(x, Variable)
 
-        Returns:
-            A list of SciPy CSC sparse matrices [D2X, D2Y].
-        """
-        if isinstance(self.args[0], Variable) and isinstance(self.args[1], Variable):
-            return {(self.args[0], self.args[1]): np.eye(self.size), 
-                    (self.args[1], self.args[0]): np.eye(self.size)}
-        if isinstance(self.args[0], Constant) and isinstance(self.args[1], Variable):
-            return self.args[1].hess
-        x = values[0]
-        y = values[1]
-        # what is the hessian of elementwise multiplication?
-        # Flatten in case inputs are not 1D
-        x = np.asarray(x).flatten(order='F')
-        y = np.asarray(y).flatten(order='F')
-        D2X = sp.diags(y, format='csc')
-        D2Y = sp.diags(x, format='csc')
-        return [D2X, D2Y]
+        if not (both_are_variables or one_is_constant or x_is_promote or y_is_promote):
+            return False
+        
+        if both_are_variables and x.id == y.id:
+            return False
+
+        return True
+
+    def _hess_vec(self, vec):
+        x = self.args[0]
+        y = self.args[1]
+        
+        # constant * atom
+        if x.is_constant(): 
+            y_hess_vec = y.hess_vec(x.value * vec)
+            return y_hess_vec
+        
+        # atom * constant
+        if y.is_constant():
+            x_hess_vec = x.hess_vec(y.value * vec)
+            return x_hess_vec
+
+        # x * y with x a scalar variable, y a vector variable
+        if not isinstance(x, Variable) and x.is_affine():
+            assert(type(x) == Promote)
+            x_var = x.args[0] # here x is a Promote because of how we canonicalize
+            return {(x_var, y): vec, (y, x_var): vec}
+        
+        # x * y with x a vector variable, y a scalar
+        if not isinstance(y, Variable) and y.is_affine():
+            assert(type(y) == Promote)
+            y_var = y.args[0] # here y is a Promote because of how we canonicalize
+            return {(x, y_var): vec, (y_var, x): vec}
+        
+        # if we arrive here both arguments are variables of the same size
+        return {(x, y): np.diag(vec), (y, x): np.diag(vec)}
 
     def graph_implementation(
         self, arg_objs, shape: Tuple[int, ...], data=None

cvxpy/atoms/elementwise/abs.py

@@ -83,3 +83,17 @@ def _grad(self, values):
         D += (values[0] > 0)
         D -= (values[0] < 0)
         return [abs.elemwise_grad_to_diag(D, rows, cols)]
+
+    def _verify_hess_vec_args(self):
+        return True
+
+    def _hess_vec(self, vec):
+        """
+        Computes the Hessian-vector product dictionary
+        for the abs atom. We assume that the argument will be a variable.
+        """
+        raise NotImplementedError("Second derivative of abs is not implemented yet.")
+        hess_dict = {}
+        var = self.args[0]
+        hess_dict[(var, var)] = np.diag(np.sign(var.value) * vec)
+        return hess_dict

cvxpy/atoms/elementwise/kl_div.py

@@ -22,6 +22,7 @@
 
 from cvxpy.atoms.elementwise.elementwise import Elementwise
 from cvxpy.constraints.constraint import Constraint
+from cvxpy.expressions.variable import Variable
 
 
 class kl_div(Elementwise):
@@ -90,6 +91,14 @@ def _grad(self, values) -> List[Optional[csc_array]]:
                                                            rows, cols)]
             return grad_list
 
+    def _verify_hess_vec_args(self):
+        return isinstance(self.args[0], Variable)
+
+    def _hess_vec(self, vec):
+        """ See the docstring of the hess_vec method of the atom class. """
+        x = self.args[0]
+        return {(x, x): np.diag(-vec / x.value)}
+    
     def _domain(self) -> List[Constraint]:
         """Returns constraints describing the domain of the node.
         """

cvxpy/atoms/elementwise/maximum.py

@@ -109,3 +109,10 @@ def _grad(self, values) -> List[Any]:
             grad_list += [maximum.elemwise_grad_to_diag(grad_vals,
                                                         rows, cols)]
         return grad_list
+
+    def _verify_hess_vec_args(self):
+        return True
+
+    def _hess_vec(self, vec):
+        """See the docstring of the hess_vec method of the atom class."""
+        raise NotImplementedError("Second derivative of maximum is not implemented yet.")

cvxpy/atoms/elementwise/minimum.py

@@ -101,3 +101,10 @@ def _grad(self, values) -> List[Any]:
             grad_list += [minimum.elemwise_grad_to_diag(grad_vals,
                                                         rows, cols)]
         return grad_list
+
+    def _verify_hess_vec_args(self):
+        return True
+
+    def _hess_vec(self, vec):
+        """ See the docstring of the hess_vec method of the atom class. """
+        raise NotImplementedError("Second derivative of minimum is not implemented yet.")

cvxpy/atoms/elementwise/trig.py

@@ -83,6 +83,18 @@ def _grad(self, values) -> List[Constraint]:
         grad_vals = np.cos(values[0])
         return [sin.elemwise_grad_to_diag(grad_vals, rows, cols)]
 
+    def _verify_hess_vec_args(self):
+        return True
+
+    def _hess_vec(self, vec):
+        """
+        Computes the Hessian-vector product dictionary
+        for the sin atom. We assume that the argument will be a variable.
+        """
+        hess_dict = {}
+        var = self.args[0]
+        hess_dict[(var, var)] = np.diag(-np.sin(var.value) * vec)
+        return hess_dict
 
 class cos(Elementwise):
     """Elementwise :math:`\\cos x`.
@@ -145,6 +157,19 @@ def _grad(self, values) -> List[Constraint]:
         cols = self.size
         grad_vals = -np.sin(values[0])
         return [cos.elemwise_grad_to_diag(grad_vals, rows, cols)]
+    
+    def _verify_hess_vec_args(self):
+        return True
+
+    def _hess_vec(self, vec):
+        """
+        Computes the Hessian-vector product dictionary
+        for the cos atom. We assume that the argument will be a variable.
+        """
+        hess_dict = {}
+        var = self.args[0]
+        hess_dict[(var, var)] = np.diag(-np.cos(var.value) * vec)
+        return hess_dict
 
 
 class tan(Elementwise):
@@ -208,3 +233,16 @@ def _grad(self, values) -> List[Constraint]:
         cols = self.size
         grad_vals = 1/np.cos(values[0])**2
         return [tan.elemwise_grad_to_diag(grad_vals, rows, cols)]
+
+    def _verify_hess_vec_args(self):
+        return True
+
+    def _hess_vec(self, vec):
+        """
+        Computes the Hessian-vector product dictionary
+        for the tan atom. We assume that the argument will be a variable.
+        """
+        hess_dict = {}
+        var = self.args[0]
+        hess_dict[(var, var)] = np.diag(2*np.tan(var.value)/np.cos(var.value)**2 * vec)
+        return hess_dict

cvxpy/atoms/pnorm.py

@@ -270,3 +270,10 @@ def _column_grad(self, value):
             nominator = np.power(value, exp)
         frac = np.divide(nominator, denominator)
         return np.reshape(frac, (frac.size, 1))
+
+    def _verify_hess_vec_args(self):
+        return True
+
+    def _hess_vec(self, vec):
+        """ See the docstring of the hess_vec method of the atom class. """
+        raise NotImplementedError("Second derivative of p-norm is not implemented yet.")

cvxpy/atoms/quad_form.py

@@ -126,13 +126,21 @@ def _grad(self, values):
         D = (P + np.conj(P.T)) @ x
         return [sp.csc_array([D.ravel(order="F")]).T]
 
-    def _hess(self, values):
+    def _verify_hess_vec_args(self):
+        return True
+
+    def _hess_vec(self, vec):
         """
-        The hessian of a quadratic form x.T @ Q @ x
-        with respect to x, is the constant matrix Q.
+        Computes the Hessian-vector product dictionary
+        for a quadratic form. We assume that the quad-form will be
+        canonicalized to w.T @ Q @ w, where w is a single variable
+        and Q is a constant matrix.
         """
-        var = self.variables()[0]
-        return {(var, var): 2 * np.array(values[1])}
+        hess_dict = {}
+        var = self.args[0]
+        Q = self.args[1]
+        hess_dict[(var, var)] = vec * 2 * Q.value
+        return hess_dict
 
     def shape_from_args(self) -> Tuple[int, ...]:
         return tuple()
@@ -174,6 +182,22 @@ def sign_from_args(self) -> Tuple[bool, bool]:
     def is_quadratic(self) -> bool:
         return True
 
+    def _verify_hess_vec_args(self):
+        return True
+
+    def _hess_vec(self, vec):
+        """
+        Computes the Hessian-vector product dictionary
+        for a quadratic form. We assume that the quad-form will be
+        canonicalized to w.T @ Q @ w, where w is a single variable
+        and Q is a constant matrix.
+        """
+        hess_dict = {}
+        var = self.args[0]
+        Q = self.args[1]
+        hess_dict[(var, var)] = vec * 2 * Q.value
+        return hess_dict
+
 
 def decomp_quad(P, cond=None, rcond=None, lower=True, check_finite: bool = True):
     """