import numpy as np
from numpy.typing import NDArray
import logging
from typing import Dict, Any
from mumott.optimization.optimizers.base_optimizer import Optimizer
from mumott.optimization.loss_functions import SquaredLoss
from mumott.methods.residual_calculators import ZHTTResidualCalculator
import tqdm
from mumott.core.hashing import list_to_hash
logger = logging.getLogger(__name__)
[docs]class ZHTTOptimizer(Optimizer):
r"""
Simple optimizer meant to be used in conjunction with :class:`ZHTTResidualCalculator`.
Cost function defined this way is non-convex, so this optimizer depends on being given a good
starting guess. Such a guess can be generated by a different model and then fitted to an
axially symmetric model.
For mode details, see the doncumentation of :class:`ZHTTResidualCalculator`.
Parameters
----------
loss_function : LossFunction
The :ref:`loss function <loss_functions>` to be minimized using this algorithm.
x0
Initial guess for solution vector. Must be the same size as
:attr:`residual_calculator.coefficients`. Defaults to :attr:`loss_function.initial_values`.
step_size : float
Step size for the gradient. A largest possible safe step-size will be estimated if none is
given.
kwargs : Dict[str, Any]
Miscellaneous options. See notes for valid entries.
Notes
-----
Valid entries in :attr:`kwargs` are
maxiter : int
Maximum number of iterations. Default value is ``20``.
"""
def __init__(self,
loss_function: SquaredLoss,
x0: NDArray[float],
step_size_parameter=None,
**kwargs: Dict[str, Any]):
super().__init__(loss_function, **kwargs)
self._options['x0'] = x0
if not isinstance(self._loss_function._residual_calculator, ZHTTResidualCalculator):
raise NotImplementedError('This optimizer requires a ZHTTResidualCalculator'
' instance for calculating the residual.')
if step_size_parameter is None:
logger.info('Since the step size has not been specified the largest safe step size will be'
'estimated. This calculation is approximate and does not take into account'
'regularization. There is therefore no guarantee of convergence.')
self._caluclate_safe_step_size_parameter()
else:
self._step_size_parameter = step_size_parameter
def _caluclate_safe_step_size_parameter(self):
""" Generate an estimated largest safe step-size for the optimization algorithm.
"""
self._step_size_parameter = self._loss_function.get_estimate_of_lifschitz_constant()
[docs] def optimize(self):
# Default parameters
opt_kwargs = dict(maxiter=20,
x0=None)
# update with value from kwargs
for k in opt_kwargs:
if k in dict(self):
opt_kwargs[k] = self[k] # Note: Base class has a __getitem__ that returns kwargs items
# Print warning for unrecognizes kwargs items
for k in dict(self):
if k not in opt_kwargs:
logger.warning(f'Unknown option {k} with value {self[k]} will be ignored.')
# prepare optimization
x = opt_kwargs['x0']
loss_function_output = self._loss_function.get_loss(x,
gradient_part='full')
# Toggle between printing an error bar or not
if not self._no_tqdm:
iterator = tqdm.tqdm(range(opt_kwargs['maxiter']))
iterator.set_description(f"Loss = {loss_function_output['loss']:.2E}")
elif self._no_tqdm:
iterator = range(opt_kwargs['maxiter'])
for ii in iterator:
# Get gradient
loss_function_output = self._loss_function.get_loss(x,
get_gradient=True,
gradient_part='full')
# Angle step size has to scale with the absolute scale of the coefficients
x = self._loss_function._residual_calculator.coefficients
step_size_scale = np.array(x[..., 0])
step_size_scale[step_size_scale < 0] = 0
step_size_scale = step_size_scale**2 + 1e-15
x[..., :-2] = x[..., :-2] - self._step_size_parameter*loss_function_output['gradient'][..., :-2]
x[..., -2:] = x[..., -2:] - self._step_size_parameter/step_size_scale[..., np.newaxis]\
* loss_function_output['gradient'][..., -2:]
if not self._no_tqdm:
iterator.set_description(f"Loss = {loss_function_output['loss']:.2E}")
loss_function_output = self._loss_function.get_loss(x,
gradient_part='full')
result = dict(x=x, loss=loss_function_output['loss'], nit=ii+1)
return dict(result)
def __hash__(self) -> int:
to_hash = [self._options, hash(self._loss_function), hash(self._step_size_parameter)]
return int(list_to_hash(to_hash), 16)