Merge pull request #3 from brunorigal/update_doc_format

Update doc format

autograd_minimize/base_wrapper.py

@@ -1,33 +1,39 @@
 from abc import ABC, abstractmethod
+
 import numpy as np
 import scipy.optimize as sopt
 
 
 class BaseWrapper(ABC):
-
     def get_input(self, input_var):
         self.input_type = type(input_var)
         assert self.input_type in [
-            dict, list, np.ndarray], 'The initial input to your optimized function should be one of dict, list or np.ndarray'
+            dict,
+            list,
+            np.ndarray,
+        ], "The initial input to your optimized function should be one of dict, list or np.ndarray"
         input_, self.shapes = self._concat(input_var)
         self.var_num = input_.shape[0]
         return input_
 
     def get_output(self, output_var):
-        assert 'shapes' in dir(
-            self), 'You must first call get input to define the tensors shapes.'
+        assert "shapes" in dir(
+            self
+        ), "You must first call get input to define the tensors shapes."
         output_var_ = self._unconcat(output_var, self.shapes)
         return output_var_
 
     def get_bounds(self, bounds):
 
         if bounds is not None:
-            if isinstance(bounds, tuple) and not (isinstance(bounds[0], tuple) or isinstance(bounds[0], sopt.Bounds)):
+            if isinstance(bounds, tuple) and not (
+                isinstance(bounds[0], tuple) or isinstance(bounds[0], sopt.Bounds)
+            ):
                 assert len(bounds) == 2
-                new_bounds = [bounds]*self.var_num
+                new_bounds = [bounds] * self.var_num
 
             elif isinstance(bounds, sopt.Bounds):
-                new_bounds = [bounds]*self.var_num
+                new_bounds = [bounds] * self.var_num
 
             elif type(bounds) in [list, tuple, np.ndarray]:
                 if self.input_type in [list, tuple]:
@@ -47,34 +53,37 @@ def get_bounds(self, bounds):
                     if k in bounds.keys():
                         new_bounds += format_bounds(bounds[k], self.shapes[k])
                     else:
-                        new_bounds += [(None, None)
-                                       ]**np.prod(self.shapes[k], dtype=np.int32)
+                        new_bounds += [(None, None)] * np.prod(
+                            self.shapes[k], dtype=np.int32
+                        )
         else:
             new_bounds = bounds
         return new_bounds
 
     def get_constraints(self, constraints, method):
-        if constraints is not None and not isinstance(constraints, sopt.LinearConstraint):
+        if constraints is not None and not isinstance(
+            constraints, sopt.LinearConstraint
+        ):
             assert isinstance(constraints, dict)
-            assert 'fun' in constraints.keys()
-            self.ctr_func = constraints['fun']
-            use_autograd = constraints.get('use_autograd', True)
-            if method in ['trust-constr']:
+            assert "fun" in constraints.keys()
+            self.ctr_func = constraints["fun"]
+            use_autograd = constraints.get("use_autograd", True)
+            if method in ["trust-constr"]:
 
                 constraints = sopt.NonlinearConstraint(
                     self._eval_ctr_func,
-                    lb=constraints.get('lb', -np.inf),
-                    ub=constraints.get('ub', np.inf),
-                    jac='2-point',
-                    keep_feasible=constraints.get('keep_feasible', False),
+                    lb=constraints.get("lb", -np.inf),
+                    ub=constraints.get("ub", np.inf),
+                    jac="2-point",
+                    keep_feasible=constraints.get("keep_feasible", False),
                 )
-            elif method in ['COBYLA', 'SLSQP']:
+            elif method in ["COBYLA", "SLSQP"]:
                 constraints = {
-                    'type': constraints.get('type', 'eq'),
-                    'fun': self._eval_ctr_func,
+                    "type": constraints.get("type", "eq"),
+                    "fun": self._eval_ctr_func,
                 }
                 if use_autograd:
-                    constraints['jac'] = self.get_ctr_jac
+                    constraints["jac"] = self.get_ctr_jac
             else:
                 raise NotImplementedError
         elif constraints is None:
@@ -153,9 +162,10 @@ def _unconcat(self, ten, shapes):
         if isinstance(shapes, dict):
             ten_vals = {}
             for k, sh in shapes.items():
-                next_ind = current_ind+np.prod(sh, dtype=np.int32)
+                next_ind = current_ind + np.prod(sh, dtype=np.int32)
                 ten_vals[k] = self._reshape(
-                    self._gather(ten, current_ind, next_ind), sh)
+                    self._gather(ten, current_ind, next_ind), sh
+                )
                 current_ind = next_ind
 
         elif isinstance(shapes, list) or isinstance(shapes, tuple):
@@ -164,9 +174,10 @@ def _unconcat(self, ten, shapes):
             else:
                 ten_vals = []
                 for sh in shapes:
-                    next_ind = current_ind+np.prod(sh, dtype=np.int32)
+                    next_ind = current_ind + np.prod(sh, dtype=np.int32)
                     ten_vals.append(
-                        self._reshape(self._gather(ten, current_ind, next_ind), sh))
+                        self._reshape(self._gather(ten, current_ind, next_ind), sh)
+                    )
 
                     current_ind = next_ind
 
@@ -191,9 +202,9 @@ def _gather(self, t, i, j):
 def format_bounds(bounds_, sh):
     if isinstance(bounds_, tuple):
         assert len(bounds_) == 2
-        return [bounds_]*np.prod(sh, dtype=np.int32)
+        return [bounds_] * np.prod(sh, dtype=np.int32)
     elif isinstance(bounds_, sopt.Bounds):
-        return [bounds_]*np.prod(sh, dtype=np.int32)
+        return [bounds_] * np.prod(sh, dtype=np.int32)
     elif isinstance(bounds_, list):
         assert np.prod(sh) == len(bounds_)
         return bounds_

autograd_minimize/scipy_minimize.py

@@ -1,121 +1,135 @@
 import scipy.optimize as sopt
 
 
-def minimize(fun, x0, backend='tf', precision='float32', method=None,
-             hvp_type=None, torch_device='cpu',
-             bounds=None, constraints=None, tol=None, callback=None, options=None):
+def minimize(
+    fun,
+    x0,
+    backend="tf",
+    precision="float32",
+    method=None,
+    hvp_type=None,
+    torch_device="cpu",
+    bounds=None,
+    constraints=None,
+    tol=None,
+    callback=None,
+    options=None,
+):
     """
     wrapper around the [minimize](https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.minimize.html)
-    function of scipy which includes an automatic computation of gradients, 
-    hessian vector product or hessian with tensorflow or torch backends. 
-
-    :param fun: function to be minimized, its signature can be a tensor, a list of tensors or a dict of tensors.
-    :type fun: tensorflow of torch function
-
-    :param x0: input to the function, it must match the signature of the function.
-    :type x0: np.ndarray, list of arrays or dict of arrays. 
-
-    :param backend: one of 'tf' or 'torch', defaults to 'tf'
-    :type backend: str, optional
-
-    :param precision: one of 'float32' or 'float64', defaults to 'float32'
-    :type precision: str, optional
-
-    :param method: method used by the optimizer, it should be one of: 
-        'Nelder-Mead', 
-        'Powell', 
-        'CG', 
-        'BFGS', 
-        'Newton-CG', 
-        'L-BFGS-B', 
-        'TNC', 
-        'COBYLA', 
-        'SLSQP', 
+    function of scipy which includes an automatic computation of gradients,
+    hessian vector product or hessian with tensorflow or torch backends.
+
+    Args:
+    * fun: function to be minimized, its signature can be a tensor, a list of tensors or a dict of tensors.
+
+    * x0: input to the function, it must match the signature of the function.
+
+    * backend: one of 'tf' or 'torch', defaults to 'tf'
+
+    * precision: one of 'float32' or 'float64', defaults to 'float32'
+
+    * method: method used by the optimizer, it should be one of:
+        'Nelder-Mead',
+        'Powell',
+        'CG',
+        'BFGS',
+        'Newton-CG',
+        'L-BFGS-B',
+        'TNC',
+        'COBYLA',
+        'SLSQP',
         'trust-constr',
         'dogleg',  # requires positive semi definite hessian
-        'trust-ncg', 
+        'trust-ncg',
         'trust-exact', # requires hessian
         'trust-krylov'
         , defaults to None
-    :type method: str, optional
 
-    :param hvp_type: type of computation scheme for the hessian vector product
+    * hvp_type: type of computation scheme for the hessian vector product
         for the torch backend it is one of hvp and vhp (vhp is faster according to the [doc](https://pytorch.org/docs/stable/autograd.html))
         for the tf backend it is one of 'forward_over_back', 'back_over_forward', 'tf_gradients_forward_over_back' and 'back_over_back'
         Some infos about the most interesting scheme are given [here](https://www.tensorflow.org/api_docs/python/tf/autodiff/ForwardAccumulator)
         , defaults to None
-    :type hvp_type: str, optional
 
-    :param torch_device: device used by torch for the gradients computation, 
+    * torch_device: device used by torch for the gradients computation,
         if the backend is not torch, this parameter is ignored, defaults to 'cpu'
-    :type torch_device: str, optional
 
-    :param bounds: Bounds on the input variables, only available for L-BFGS-B, TNC, SLSQP, Powell, and trust-constr methods.
-        It can be: 
-        * a tuple (min, max), None indicates no bounds, in this case the same bound is applied to all variables. 
-        * An instance of the [Bounds](https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.Bounds.html#scipy.optimize.Bounds) class, in this case the same bound is applied to all variables. 
+    * bounds: Bounds on the input variables, only available for L-BFGS-B, TNC, SLSQP, Powell, and trust-constr methods.
+        It can be:
+        * a tuple (min, max), None indicates no bounds, in this case the same bound is applied to all variables.
+        * An instance of the [Bounds](https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.Bounds.html#scipy.optimize.Bounds) class, in this case the same bound is applied to all variables.
         * A numpy array of bounds (if the optimized function has a single numpy array as input)
-        * A list or dict of bounds with the same format as the optimized function signature. 
+        * A list or dict of bounds with the same format as the optimized function signature.
         , defaults to None
-    :type bounds: tuple, list, dict or np.ndarray, optional
 
-    :param constraints: It has to be a dict with the following keys:
+    * constraints: It has to be a dict with the following keys:
         * fun: a callable computing the constraint function
         * lb and ub: the lower and upper bounds, if equal, the constraint is an inequality, use np.inf if there is no upper bound. Only used if method is trust-constr.
         * type: 'eq' or 'ineq' only used if method is one of COBYLA, SLSQP.
         * keep_feasible: see [here](https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.NonlinearConstraint.html#scipy.optimize.NonlinearConstraint)
         , defaults to None
-    :type constraints: dict, optional
 
-    :param tol: Tolerance for termination, defaults to None
-    :type tol: float, optional
+    * tol: Tolerance for termination, defaults to None
 
-    :param callback: Called after each iteration, defaults to None
-    :type callback: callable, optional
+    * callback: Called after each iteration, defaults to None
 
-    :param options: solver options, defaults to None
-    :type options: dict, optional
+    * options: solver options, defaults to None
 
-    :return: dict of optimization results
-    :rtype: dict
     """
 
-    if backend == 'tf':
+    if backend == "tf":
         from .tf_wrapper import TfWrapper
+
         wrapper = TfWrapper(fun, precision=precision, hvp_type=hvp_type)
-    elif backend == 'torch':
+    elif backend == "torch":
         from .torch_wrapper import TorchWrapper
-        wrapper = TorchWrapper(fun, precision=precision,
-                               hvp_type=hvp_type, device=torch_device)
+
+        wrapper = TorchWrapper(
+            fun, precision=precision, hvp_type=hvp_type, device=torch_device
+        )
     else:
         raise NotImplementedError
 
     if bounds is not None:
-        assert method in [None,  'L-BFGS-B', 'TNC', 'SLSQP', 'Powell',
-                          'trust-constr'], 'bounds are only available for L-BFGS-B, TNC, SLSQP, Powell, trust-constr'
+        assert method in [
+            None,
+            "L-BFGS-B",
+            "TNC",
+            "SLSQP",
+            "Powell",
+            "trust-constr",
+        ], "bounds are only available for L-BFGS-B, TNC, SLSQP, Powell, trust-constr"
 
     if constraints is not None:
         assert method in [
-            'COBYLA', 'SLSQP', 'trust-constr'], 'Constraints are only available for COBYLA, SLSQP and trust-constr'
-
-    optim_res = sopt.minimize(wrapper.get_value_and_grad,
-                              wrapper.get_input(x0), method=method, jac=True,
-                              hessp=wrapper.get_hvp if method in ['Newton-CG', 'trust-ncg',
-                                                                  'trust-krylov', 'trust-constr']
-                              else None,
-                              hess=wrapper.get_hess if method in [
-                                  'dogleg', 'trust-exact'] else None,
-                              bounds=wrapper.get_bounds(bounds),
-                              constraints=wrapper.get_constraints(
-                                  constraints, method),
-                              tol=tol, callback=callback, options=options)
+            "COBYLA",
+            "SLSQP",
+            "trust-constr",
+        ], "Constraints are only available for COBYLA, SLSQP and trust-constr"
+
+    optim_res = sopt.minimize(
+        wrapper.get_value_and_grad,
+        wrapper.get_input(x0),
+        method=method,
+        jac=True,
+        hessp=wrapper.get_hvp
+        if method in ["Newton-CG", "trust-ncg", "trust-krylov", "trust-constr"]
+        else None,
+        hess=wrapper.get_hess if method in ["dogleg", "trust-exact"] else None,
+        bounds=wrapper.get_bounds(bounds),
+        constraints=wrapper.get_constraints(constraints, method),
+        tol=tol,
+        callback=callback,
+        options=options,
+    )
 
     optim_res.x = wrapper.get_output(optim_res.x)
 
-    if 'jac' in optim_res.keys() and len(optim_res.jac) > 0:
+    if "jac" in optim_res.keys() and len(optim_res.jac) > 0:
         try:
             optim_res.jac = wrapper.get_output(optim_res.jac[0])
-        except:
+        except Exception:
             pass
 
     return optim_res