"""Optimizers implementations."""
import logging
from abc import ABC, abstractmethod
import scipy
from scipy.optimize import OptimizeResult
import geomstats.backend as gs
from geomstats.exceptions import AutodiffNotImplementedError
from geomstats.numerics._common import result_to_backend_type
[docs]
class ScipyMinimize:
"""Wrapper for scipy.optimize.minimize.
Only `autodiff_jac` and `autodiff_hess` differ from scipy: if True, then
automatic differentiation is used to compute jacobian and/or hessian.
"""
def __init__(
self,
method="L-BFGS-B",
autodiff_jac=False,
autodiff_hess=False,
bounds=None,
constraints=(),
tol=None,
callback=None,
options=None,
save_result=False,
):
if (autodiff_jac or autodiff_hess) and not gs.has_autodiff():
raise AutodiffNotImplementedError(
"Minimization with 'autodiff' requires automatic differentiation."
"Change backend via the command "
"export GEOMSTATS_BACKEND=pytorch in a terminal"
)
self.method = method
self.autodiff_jac = autodiff_jac
self.autodiff_hess = autodiff_hess
self.bounds = bounds
self.constraints = constraints
self.tol = tol
self.callback = callback
self.options = options
self.save_result = save_result
self.result_ = None
def _handle_jac(self, fun, fun_jac):
if fun_jac is not None:
fun_ = lambda x: fun(gs.from_numpy(x))
fun_jac_ = fun_jac
if callable(fun_jac):
fun_jac_ = lambda x: fun_jac(gs.from_numpy(x))
return fun_, fun_jac_
if self.autodiff_jac:
jac = True
fun_ = lambda x: gs.autodiff.value_and_grad(fun)(gs.from_numpy(x))
else:
jac = fun_jac
fun_ = lambda x: fun(gs.from_numpy(x))
return fun_, jac
def _handle_hess(self, fun, fun_hess):
if fun_hess is not None or not self.autodiff_hess:
fun_hess_ = fun_hess
if callable(fun_hess):
fun_hess_ = lambda x: fun_hess(gs.from_numpy(x))
return fun_hess_
return lambda x: gs.autodiff.hessian(fun)(gs.from_numpy(x))
[docs]
def minimize(self, fun, x0, fun_jac=None, fun_hess=None, hessp=None):
"""Minimize objective function.
Parameters
----------
fun : callable
The objective function to be minimized.
x0 : array-like
Initial guess.
fun_jac : callable
Jacobian of fun.
fun_hess : callable
Hessian of fun.
hessp : callable
"""
fun_, jac = self._handle_jac(fun, fun_jac)
hess = self._handle_hess(fun, fun_hess)
result = scipy.optimize.minimize(
fun_,
x0,
method=self.method,
jac=jac,
hess=hess,
hessp=hessp,
bounds=self.bounds,
tol=self.tol,
constraints=self.constraints,
callback=self.callback,
options=self.options,
)
result = result_to_backend_type(result)
if not result.success:
logging.warning(result.message)
if self.save_result:
self.result_ = result
return result
[docs]
class RootFinder(ABC):
"""Find a root of a vector-valued function."""
[docs]
@abstractmethod
def root(self, fun, x0, fun_jac=None):
"""Find a root of a vector-valued function.
Parameters
----------
fun : callable
Vector-valued function.
x0 : array-like
Initial guess.
fun_jac : callable
Ignored if None.
Returns
-------
res : OptimizeResult
"""
[docs]
class ScipyRoot(RootFinder):
"""Wrapper for scipy.optimize.root.
Only `autodiff_jac` differs from scipy: if True, then
automatic differentiation is used to compute jacobian.
"""
def __init__(
self,
method="hybr",
autodiff_jac=False,
tol=None,
callback=None,
options=None,
save_result=False,
):
if (autodiff_jac) and not gs.has_autodiff():
raise AutodiffNotImplementedError(
"Root finding with 'autodiff' requires automatic differentiation."
"Change backend via the command "
"export GEOMSTATS_BACKEND=pytorch in a terminal"
)
self.method = method
self.autodiff_jac = autodiff_jac
self.tol = tol
self.callback = callback
self.options = options
self.save_result = save_result
self.result_ = None
def _handle_jac(self, fun, fun_jac):
if fun_jac is not None:
fun_ = lambda x: fun(gs.from_numpy(x))
fun_jac_ = fun_jac
if callable(fun_jac):
fun_jac_ = lambda x: fun_jac(gs.from_numpy(x))
return fun_, fun_jac_
if self.autodiff_jac:
jac = True
fun_ = lambda x: gs.autodiff.value_and_jacobian(fun)(gs.from_numpy(x))
else:
jac = fun_jac
fun_ = lambda x: fun(gs.from_numpy(x))
return fun_, jac
[docs]
def root(self, fun, x0, fun_jac=None):
"""Find a root of a vector-valued function.
Parameters
----------
fun : callable
Vector-valued function.
x0 : array-like
Initial guess.
fun_jac : callable
Jacobian of fun. Ignored if None.
Returns
-------
res : OptimizeResult
"""
fun_, fun_jac_ = self._handle_jac(fun, fun_jac)
result = scipy.optimize.root(
fun_,
x0,
method=self.method,
jac=fun_jac_,
tol=self.tol,
callback=self.callback,
options=self.options,
)
result = result_to_backend_type(result)
if not result.success:
logging.warning(result.message)
if self.save_result:
self.result_ = result
return result
[docs]
class NewtonMethod(RootFinder):
"""Find a root of a vector-valued function with Newton's method.
Check https://en.wikipedia.org/wiki/Newton%27s_method_in_optimization
for details.
Parameters
----------
atol : float
Tolerance to check algorithm convergence.
max_iter : int
Maximum iterations.
"""
def __init__(self, atol=gs.atol, max_iter=100):
self.atol = atol
self.max_iter = max_iter
[docs]
def root(self, fun, x0, fun_jac):
"""Find a root of a vector-valued function.
Parameters
----------
fun : callable
Vector-valued function.
x0 : array-like
Initial guess.
fun_jac : callable
Jacobian of fun.
"""
xk = x0
message = "The solution converged."
status = 1
for it in range(self.max_iter):
fun_xk = fun(xk)
if gs.linalg.norm(fun_xk) <= self.atol:
break
y = gs.linalg.solve(fun_jac(xk), fun_xk)
xk = xk - y
else:
message = f"Maximum number of iterations {self.max_iter} reached. Results may be inaccurate"
status = 0
result = OptimizeResult(
x=xk,
success=(status == 1),
status=status,
method="newton",
message=message,
nfev=it + 1,
njac=it,
)
if not result.success:
logging.warning(result.message)
return result