Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Update doc format #3

Merged
merged 5 commits into from
Feb 8, 2023
Merged

Update doc format #3

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
560ed6d
wrong ** instead of * for list creation
Feb 8, 2023
6e5dd92
black
Feb 8, 2023
60071b5
isort
Feb 8, 2023
983befb
update doc
Feb 8, 2023
5ee1bee
adds requirements
Feb 8, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
67 changes: 39 additions & 28 deletions autograd_minimize/base_wrapper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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]:
Expand All @@ -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:
Expand Down Expand Up @@ -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):
Expand All @@ -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

Expand All @@ -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_
Expand Down
156 changes: 85 additions & 71 deletions autograd_minimize/scipy_minimize.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Loading