Inverse

inverse

This module provides classes to perform an adjoint inverse optimisation and checkpoint intermediate results. Users instantiate the LinMoreOptimiser class by providing relevant parameters and call the run method to perform the optimisation.

ROLSolver(problem, parameters, inner_product='L2', vector_class=pyadjoint_rol.ROLVector, vector_args=[])

Bases: ROLSolver

A ROLSolver that may use a class other than ROLVector to hold its vectors.

In the non-checkpointing case, this reduces down to the original ROLSolver class. However, if ROL checkpointing is enabled, every vector within the solver needs to be a CheckpointedROLVector. We can achieve this by overwriting the base self.rolvector member.

Initialises the solver instance.

Parameters:

Name	Type	Description	Default
problem	MinimizationProblem	PyAdjoint minimisation problem to solve.	required
parameters	dict	Dictionary containing the parameters governing ROL.	required
inner_product	str	String representing the inner product to use for vector operations.	'L2'
vector_class		PyAdjoint ROL vector representing the underlying vector class.	ROLVector
vector_args	list	List of arguments for initialisation of the vector class.	[]

Source code in g-adopt/gadopt/inverse.py

def __init__(
    self,
    problem: MinimizationProblem,
    parameters: dict,
    inner_product: str = "L2",
    vector_class=pyadjoint_rol.ROLVector,
    vector_args: list = [],
):
    """Initialises the solver instance.

    Args:
      problem:
        PyAdjoint minimisation problem to solve.
      parameters:
        Dictionary containing the parameters governing ROL.
      inner_product:
        String representing the inner product to use for vector operations.
      vector_class:
        PyAdjoint ROL vector representing the underlying vector class.
      vector_args:
        List of arguments for initialisation of the vector class.
    """
    super().__init__(problem, parameters)

    x = [p.tape_value() for p in self.problem.reduced_functional.controls]
    self.rolvector = vector_class(x, *vector_args, inner_product=inner_product)

    # we need to recompute these with the new rolvector instance
    self.bounds = self.__get_bounds()
    self.constraints = self.__get_constraints()

LinMoreOptimiser(problem, parameters, checkpoint_dir=None, auto_checkpoint=True)

The management class for Lin-More trust region optimisation using ROL.

This class sets up ROL to use the Lin-More trust region method, with a limited-memory BFGS secant for determining the steps. A pyadjoint problem has to be set up first, containing the optimisation functional and other constraints (like bounds).

This optimiser also supports checkpointing ROL's state, to allow resumption of a previous optimisation without having to refill the L-BFGS memory. The underlying objects will be configured for checkpointing if checkpoint_dir is specified, and optionally the automatic checkpointing each iteration can be disabled by the auto_checkpoint parameter.

Parameters:

Name	Type	Description	Default
problem	MinimizationProblem	The actual problem to solve.	required
parameters	dict	A dictionary containing the parameters governing ROL.	required
checkpoint_dir	Optional[str]	A path to hold any checkpoints that are saved.	None
auto_checkpoint	Optional[bool]	Whether to automatically checkpoint each iteration.	True

Source code in g-adopt/gadopt/inverse.py

def __init__(self, problem, parameters, checkpoint_dir=None, auto_checkpoint=True):
    """The management class for Lin-More trust region optimisation using ROL.

    This class sets up ROL to use the Lin-More trust region method, with a limited-memory
    BFGS secant for determining the steps. A pyadjoint problem has to be set up first,
    containing the optimisation functional and other constraints (like bounds).

    This optimiser also supports checkpointing ROL's state, to allow resumption of
    a previous optimisation without having to refill the L-BFGS memory. The underlying
    objects will be configured for checkpointing if `checkpoint_dir` is specified,
    and optionally the automatic checkpointing each iteration can be disabled by the
    `auto_checkpoint` parameter.

    Params:
        problem (pyadjoint.MinimizationProblem): The actual problem to solve.
        parameters (dict): A dictionary containing the parameters governing ROL.
        checkpoint_dir (Optional[str]): A path to hold any checkpoints that are saved.
        auto_checkpoint (Optional[bool]): Whether to automatically checkpoint each iteration.
    """

    self.iteration = -1

    solver_kwargs = {}
    if checkpoint_dir is not None:
        self._base_checkpoint_dir = Path(checkpoint_dir)
        self._ensure_checkpoint_dir()

        self._mesh = problem.reduced_functional.controls[0].control.function_space().mesh()
        solver_kwargs["vector_class"] = CheckpointedROLVector
        solver_kwargs["vector_args"] = [self]

        self.auto_checkpoint = auto_checkpoint
    else:
        self._base_checkpoint_dir = None
        self.auto_checkpoint = False

    self.rol_solver = ROLSolver(problem, parameters, inner_product="L2", **solver_kwargs)
    self.rol_parameters = ROL.ParameterList(parameters, "Parameters")

    try:
        self.rol_secant = ROL.lBFGS(parameters["General"]["Secant"]["Maximum Storage"])
    except KeyError:
        # Use the default storage value
        self.rol_secant = ROL.lBFGS()

    self.rol_algorithm = ROL.LinMoreAlgorithm(self.rol_parameters, self.rol_secant)
    self.callbacks = []

    self._add_statustest()

checkpoint()

Checkpoints the current ROL state to disk.

Source code in g-adopt/gadopt/inverse.py

def checkpoint(self):
    """Checkpoints the current ROL state to disk."""

    ROL.serialise_algorithm(self.rol_algorithm, MPI.COMM_WORLD.rank, str(self.checkpoint_dir))

    checkpoint_path = self.checkpoint_dir / "solution_checkpoint.h5"
    with CheckpointFile(str(checkpoint_path), "w") as f:
        for i, func in enumerate(self.rol_solver.rolvector.dat):
            f.save_function(func, name=f"dat_{i}")

restore(iteration=None)

Restores the ROL state from disk.

The last stored iteration in checkpoint_dir is used unless a given iteration is specifed.

Source code in g-adopt/gadopt/inverse.py

def restore(self, iteration=None):
    """Restores the ROL state from disk.

    The last stored iteration in `checkpoint_dir` is used unless a given iteration
    is specifed.
    """
    if iteration is not None:
        self.iteration = iteration
    else:
        stored_iterations = [int(path.parts[-1]) for path in self._base_checkpoint_dir.glob('*[0-9]/')]
        self.iteration = sorted(stored_iterations)[-1]

    self.rol_algorithm = ROL.load_algorithm(MPI.COMM_WORLD.rank, str(self.checkpoint_dir))
    self._add_statustest()

    self.rol_solver.rolvector.checkpoint_path = self.checkpoint_dir / "solution_checkpoint.h5"
    self.rol_solver.rolvector.load(self._mesh)

    # ROL algorithms run in a loop like `while (statusTest()) { ... }`
    # so we will double up on saving the restored iteration
    # by rolling back the iteration counter, we make sure we overwrite the checkpoint
    # we just restored, to keep the ROL iteration count, and our checkpoint iteration
    # count in sync
    self.iteration -= 1

    # The various ROLVector objects can load all their metadata, but can't actually
    # restore from the Firedrake checkpoint. They register themselves, so we can access
    # them through a flat list.
    vec = self.rol_solver.rolvector.dat
    for v in _vector_registry:
        x = [p.copy(deepcopy=True) for p in vec]
        v.dat = x
        v.load(self._mesh)
        v._optimiser = self

    _vector_registry.clear()

run()

Runs the actual ROL optimisation.

This will continue until the status test flags the optimisation to complete.

Source code in g-adopt/gadopt/inverse.py

def run(self):
    """Runs the actual ROL optimisation.

    This will continue until the status test flags the optimisation to complete.
    """

    with stop_annotating():
        self.rol_algorithm.run(
            self.rol_solver.rolvector,
            self.rol_solver.rolobjective,
            self.rol_solver.bounds,
        )

add_callback(callback)

Adds a callback to run after every optimisation iteration.

Source code in g-adopt/gadopt/inverse.py

def add_callback(self, callback):
    """Adds a callback to run after every optimisation iteration."""

    self.callbacks.append(callback)

CheckpointedROLVector(dat, optimiser, inner_product='L2')

Bases: ROLVector

An extension of ROLVector that supports checkpointing to disk.

The ROLVector itself is the Python-side implementation of the ROL.Vector interface; it defines all the operations ROL may perform on vectors (e.g., scaling, addition) and links ROL to the underlying Firedrake vectors.

When the serialisation library hits a ROL.Vector on the C++ side, it will pickle this object, so we provide implementations of __getstate__ and __setstate__ that will correctly participate in the serialisation pipeline.

Initialises the checkpointed ROL vector instance.

Parameters:

Name	Type	Description	Default
dat	Function	Firedrake function.	required
optimiser	LinMoreOptimiser	The managing optimiser for controlling checkpointing paths.	required
inner_product	str	String representing the inner product to use for vector operations.	'L2'

Source code in g-adopt/gadopt/inverse.py

def __init__(
    self,
    dat: Function,
    optimiser: LinMoreOptimiser,
    inner_product: str = "L2",
):
    """Initialises the checkpointed ROL vector instance.

    Args:
      dat:
        Firedrake function.
      optimiser (LinMoreOptimiser):
        The managing optimiser for controlling checkpointing paths.
      inner_product:
        String representing the inner product to use for vector operations.
    """
    super().__init__(dat, inner_product)

    self.comm = dat[0].comm
    self._optimiser = optimiser

save(checkpoint_path)

Checkpoints the data within this vector to disk.

Called when this object is pickled as part of the ROL state serialisation.

Source code in g-adopt/gadopt/inverse.py

def save(self, checkpoint_path):
    """Checkpoints the data within this vector to disk.

    Called when this object is pickled as part of the ROL
    state serialisation.
    """

    with CheckpointFile(str(checkpoint_path), "w") as f:
        for i, func in enumerate(self.dat):
            f.save_function(func, name=f"dat_{i}")

load(mesh)

Loads the checkpointed data for this vector from disk.

Called by the parent Optimiser after the ROL state has been deserialised. The pickling routine will register this vector within the registry.

Source code in g-adopt/gadopt/inverse.py

def load(self, mesh):
    """Loads the checkpointed data for this vector from disk.

    Called by the parent Optimiser after the ROL state has
    been deserialised. The pickling routine will register
    this vector within the registry.
    """

    with CheckpointFile(str(self.checkpoint_path), "r") as f:
        for i in range(len(self.dat)):
            self.dat[i] = f.load_function(mesh, name=f"dat_{i}")

__setstate__(state)

Sets the state from the result of unpickling.

This happens during the restoration of a checkpoint. self.dat needs to be separately set, then self.load() can be called.

Source code in g-adopt/gadopt/inverse.py

def __setstate__(self, state):
    """Sets the state from the result of unpickling.

    This happens during the restoration of a checkpoint. self.dat needs to be
    separately set, then self.load() can be called.
    """

    # initialise C++ state
    super().__init__(state)
    self.checkpoint_path, self.inner_product = state

    _vector_registry.append(self)

__getstate__()

Returns a state tuple suitable for pickling

Source code in g-adopt/gadopt/inverse.py

def __getstate__(self):
    """Returns a state tuple suitable for pickling"""

    checkpoint_id = firedrake.utils._new_uid(self.comm)
    checkpoint_filename = f"vector_checkpoint_{checkpoint_id}.h5"
    checkpoint_path = self._optimiser.checkpoint_dir / checkpoint_filename
    self.save(checkpoint_path)

    return (checkpoint_path, self.inner_product)