Source code for mumott.optimization.optimizers.zonal_harmonics_optimizer

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:'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)