otf.optim.optimizer

Optimization algorithms for estimating unknown system parameters.

This module provides concrete implementations of BaseOptimizer used to update the cs parameters of a BaseSystem instance, including simple gradient descent, Levenberg–Marquardt variants, and an adapter for Optax optimizers.

Classes:

Name	Description
DummyOptimizer	Optimizer that performs no parameter updates (useful for testing).
GradientDescent	Simple gradient-descent optimizer.
LevenbergMarquardt	Levenberg–Marquardt optimizer using sensitivity-based gradients.
OptaxWrapper	Adapter that wraps an Optax optimizer as a BaseOptimizer.
WeightedLevenbergMarquardt	Weighted Levenberg–Marquardt optimizer (Gauss–Newton variant).

DummyOptimizer

DummyOptimizer(
    system: BaseSystem,
    gradient_computer: GradientComputer | None = None,
)

Bases: BaseOptimizer

Optimizer that performs no parameter updates (useful for testing).

Initialize the dummy optimizer.

Parameters:

Name	Type	Description	Default
system	BaseSystem	Target BaseSystem instance.	required
gradient_computer	GradientComputer | None	Optional GradientComputer.	None

Source code in src/otf/optim/optimizer.py

def __init__(
    self,
    system: BaseSystem,
    gradient_computer: gradient.GradientComputer | None = None,
):
    """Initialize the dummy optimizer.

    Parameters
    ----------
    system
        Target `BaseSystem` instance.
    gradient_computer
        Optional `GradientComputer`.
    """
    super().__init__(system, gradient_computer)

GradientDescent

GradientDescent(
    system: BaseSystem,
    learning_rate: float = 0.0001,
    gradient_computer: GradientComputer | None = None,
)

Bases: BaseOptimizer

Simple gradient-descent optimizer.

Create a gradient-descent optimizer.

Parameters:

Name	Type	Description	Default
learning_rate	float	Scalar learning rate used to scale the negative gradient.	0.0001

Source code in src/otf/optim/optimizer.py

def __init__(
    self,
    system: BaseSystem,
    learning_rate: float = 1e-4,
    gradient_computer: gradient.GradientComputer | None = None,
):
    """Create a gradient-descent optimizer.

    Parameters
    ----------
    learning_rate
        Scalar learning rate used to scale the negative gradient.
    """
    super().__init__(system, gradient_computer)
    self.learning_rate = learning_rate

LevenbergMarquardt

LevenbergMarquardt(
    system: BaseSystem,
    learning_rate: float = 0.001,
    lam: float = 0.01,
    gradient_computer: SensitivityGradient | None = None,
)

Bases: BaseOptimizer

Levenberg–Marquardt optimizer using sensitivity-based gradients.

Levenberg–Marquardt optimizer using sensitivity-based gradients.

This implementation requires a SensitivityGradient instance and is currently implemented only for the UpdateOption.last_state update method of the gradient computer.

Parameters:

Name	Type	Description	Default
learning_rate	float	Scalar multiplier applied to the computed step.	0.001
lam	float	Levenberg–Marquardt damping parameter.	0.01

Source code in src/otf/optim/optimizer.py

def __init__(
    self,
    system: BaseSystem,
    learning_rate: float = 1e-3,
    lam: float = 1e-2,
    gradient_computer: gradient.SensitivityGradient | None = None,
):
    """Levenberg–Marquardt optimizer using sensitivity-based gradients.

    This implementation requires a `SensitivityGradient` instance and is
    currently implemented only for the `UpdateOption.last_state` update
    method of the gradient computer.

    Parameters
    ----------
    learning_rate
        Scalar multiplier applied to the computed step.
    lam
        Levenberg–Marquardt damping parameter.
    """
    if not isinstance(gradient_computer, gradient.SensitivityGradient):
        raise NotImplementedError(
            "not yet implemented for adjoint-based gradient computation"
        )
    if gradient_computer.update_option is not (
        gradient.sensitivity.UpdateOption.last_state
    ):
        raise NotImplementedError(
            "currently implemented only for last state gradient computation"
        )

    super().__init__(system, gradient_computer)
    self.learning_rate = learning_rate
    self.lam = lam

OptaxWrapper

OptaxWrapper(
    system: BaseSystem,
    optimizer: GradientTransformationExtraArgs,
    gradient_computer: GradientComputer | None = None,
)

Bases: BaseOptimizer

Adapter that wraps an Optax optimizer as a BaseOptimizer.

Wrap a given Optax optimizer.

Parameters:

Name	Type	Description	Default
optimizer	GradientTransformationExtraArgs	Instance of optax.GradientTransformationExtraArgs For example, optax.adam(learning_rate=1e-1).	required

Source code in src/otf/optim/optimizer.py

def __init__(
    self,
    system: BaseSystem,
    optimizer: optax.GradientTransformationExtraArgs,
    gradient_computer: gradient.GradientComputer | None = None,
):
    """Wrap a given Optax optimizer.

    Parameters
    ----------
    optimizer
        Instance of `optax.GradientTransformationExtraArgs` For example,
        `optax.adam(learning_rate=1e-1)`.
    """
    super().__init__(system, gradient_computer)
    self.optimizer = optimizer
    self.opt_state = self.optimizer.init(system.cs)

WeightedLevenbergMarquardt

WeightedLevenbergMarquardt(
    system: BaseSystem,
    learning_rate: float = 0.001,
    lam: float = 0.01,
    gradient_computer: GradientComputer | None = None,
)

Bases: BaseOptimizer

Weighted Levenberg–Marquardt optimizer (Gauss–Newton variant).

Perform a weighted version of the Levenberg–Marquardt modification of Gauss–Newton.

Parameters:

Name	Type	Description	Default
learning_rate	float	The learning rate (scalar by which to multiply the step)	0.001
lam	float	Levenberg–Marquardt parameter	0.01

Source code in src/otf/optim/optimizer.py

def __init__(
    self,
    system: BaseSystem,
    learning_rate: float = 1e-3,
    lam: float = 1e-2,
    gradient_computer: gradient.GradientComputer | None = None,
):
    """Perform a weighted version of the Levenberg–Marquardt modification of
    Gauss–Newton.

    Parameters
    ----------
    learning_rate
        The learning rate (scalar by which to multiply the step)
    lam
        Levenberg–Marquardt parameter
    """
    super().__init__(system, gradient_computer)
    self.learning_rate = learning_rate
    self.lam = lam