:orphan: :py:mod:`tlm_adjoint.checkpointing` =================================== .. py:module:: tlm_adjoint.checkpointing Module Contents --------------- .. py:class:: CheckpointStorage(*, store_ics, store_data) A buffer for forward restart data, and a cache for non-linear dependency data. Contains three types of data: 1. References: Dependencies which are stored by reference. Variables `x` for which `var_is_static(x)` is `True` are stored by reference. 2. Forward restart / initial condition data: Dependencies which are used to restart and advance the forward calculation. 3. Non-linear dependency data: Non-linear dependencies of the forward which are used to advance the adjoint. These may overlap -- for example non-linear dependency data may be stored by reference. Non-linear dependency data has an associated key `(n, i)`, where `n` is an :class:`int` indicating the block index and `i` is an :class:`int` indicating the equation index within that block. Non-linear-dependency data for an :class:`.Equation` can be accessed via, e.g. .. code-block:: python nl_deps = cp[(n, i)] where `cp` is a :class:`.CheckpointStorage`. Here `nl_deps` is a :class:`tuple` of variables storing values associated with `eq.nonlinear_dependencies()`, for :class:`.Equation` `i` in block `n`. :arg store_ics: Whether to enable storage of forward restart data. :arg store_data: Whether to enable storage of non-linear dependency data. .. !! processed by numpydoc !! .. py:property:: store_ics Whether storage of forward restart data is enabled. .. !! processed by numpydoc !! .. py:property:: store_data Whether storage of non-linear dependency data is enabled. .. !! processed by numpydoc !! .. py:method:: configure(*, store_ics, store_data) Enable or disable storage of forward restart and non-linear dependency data. :arg store_ics: Whether storage of forward restart data should be enabled (`store_ics=True`) or disabled (`store_ics=False`). :arg store_data: Whether storage of non-linear dependency data should be enabled (`store_data=True`) or disabled (`store_data=False`). .. !! processed by numpydoc !! .. py:method:: clear(*, clear_ics=True, clear_data=True, clear_refs=False) Clear stored data. :arg clear_ics: Whether forward restart data should be cleared. :arg clear_data: Whether non-linear dependency data should be cleared. :arg clear_refs: Whether references should be cleared. .. !! processed by numpydoc !! .. py:method:: initial_condition(x, *, copy=True) Return the forward restart value associated with a variable `x`. :arg x: The variable, or the :class:`int` variable ID, for which the forward restart value should be returned. :arg copy: If `True` then a copy of the stored value is returned. If `False` then an internal variable storing the value is returned. :returns: A variable containing the forward restart value for `x`. .. !! processed by numpydoc !! .. py:method:: initial_conditions(*, cp=True, refs=False, copy=True) Access stored forward restart data. :arg cp: Whether to include forward restart data that is stored by value. :arg refs: Whether to include data that is stored by reference. May include non-linear dependency data that is not forward restart data. :arg copy: If `True` then a copy of the stored data is returned. If `False` then internal variables storing the data are returned. :returns: A :class:`.VariableStateLockDictionary`, with items `(x_id: x_value)`, where `x_id` is the :class:`int` variable ID and `x_value` is a variable storing the data. .. !! processed by numpydoc !! .. py:method:: add_initial_condition(x, value=None) Store forward restart data / an initial condition dependency. :arg x: The initial condition dependency variable. :arg value: A variable defining the initial condition dependency value. `x` is used if not supplied. .. !! processed by numpydoc !! .. py:method:: update_keys(n, i, eq) The :class:`.CheckpointStorage` keeps an internal map from forward variables to equations in which values for those variables are computed. Keys are updated automatically as needed. This method allows keys to be updated manually. :arg n: The :class:`int` index of the block. :arg i: The :class:`int` index of the equation. :arg eq: An :class:`.Equation`, equation `i` in block `n`. .. !! processed by numpydoc !! .. py:method:: add_equation(n, i, eq, *, deps=None, nl_deps=None) Store checkpoint data associated with an equation. :arg n: The :class:`int` index of the block. :arg i: The :class:`int` index of the equation. :arg eq: An :class:`.Equation`, equation `i` in block `n`. :arg deps: Equation dependency values. `eq.dependencies()` is used if not supplied. :arg nl_deps: Equation non-linear dependency values. Extracted from `deps` if not supplied. .. !! processed by numpydoc !! .. py:method:: add_equation_data(n, i, eq, *, nl_deps=None) Store checkpoint data associated with an equation. As :meth:`.CheckpointStorage.add_equation`, but adds only *non-linear* dependency data. :arg n: The :class:`int` index of the block. :arg i: The :class:`int` index of the equation. :arg eq: An :class:`.Equation`, equation `i` in block `n`. :arg nl_deps: Equation non-linear dependency values. `eq.nonlinear_dependencies()` is used if not supplied. .. !! processed by numpydoc !! .. py:method:: checkpoint_data(*, ics=True, data=True, copy=True) Extract checkpoint data. :arg ics: Whether to extract forward restart data. :arg data: Whether to extract non-linear dependency data. :arg copy: If `True` then a copy of the stored data is returned. If `False` then internal variables storing the data are returned. :returns: A :class:`tuple` `(cp, data, storage)`. Elements of this :class:`tuple` are as for the three arguments for the :meth:`.CheckpointStorage.update` method, and here `storage` is a :class:`.VariableStateLockDictionary`. .. !! processed by numpydoc !! .. py:method:: update(cp, data, storage, *, copy=True) Update the :class:`.CheckpointStorage` using the provided checkpoint data. Used to update the :class:`.CheckpointStorage` from loaded data. Note that the :class:`.CheckpointStorage` is *not* cleared prior to updating using the provided data. :arg cp: A :class:`tuple` of keys. Forward restart data is defined by `(storage[key] for key in cp)`. :arg data: A :class:`dict`. Items are `((n, i), keys)`, indicating that non-linear dependency data for equation `i` in block `n` is `(storage[key] for key in keys)`. :arg storage: The stored data. A :class:`Mapping` with items `((x_id, x_indices), x_value)`. `x_id` is the :class:`int` ID for a variable whose value `x_value` is stored. `x_indices` is either `None`, if the variable value has not been computed by solving equations with forward restart data storage enabled, or a tuple `(n, i, m)` indicating that the variable value was computed as component `m` of the solution to equation `i` in block `n`. :arg copy: Whether the values in `storage` should be copied when being stored in the :class:`.CheckpointStorage`. .. !! processed by numpydoc !! .. py:class:: ReplayStorage(blocks, N0, N1, *, transpose_deps=None) Storage used when solving forward equations. A value for a forward variable can be accessed via .. code-block:: python x_value = replay_storage[x] and set via .. code-block:: python replay_storage[x] = x_value where here `x` is either a variable of an :class:`int` variable ID. Containment can also be tested, .. code-block:: python if x in replay_storage: [...] :arg blocks: A :class:`Sequence` or :class:`Mapping`, whose elements or values are :class:`Sequence` objects containing :class:`.Equation` objects. Forward equations. :arg N0: An :class:`int`. `(blocks[n] for n in range(N0, N1))` defines the forward equations which will be solved. :arg N1: An :class:`int`. `(blocks[n] for n in range(N0, N1))` defines the forward equations which will be solved. :arg transpose_deps: A :class:`.TransposeComputationalGraph`. If supplied then an activity analysis is applied. .. !! processed by numpydoc !! .. py:method:: is_active(n, i) Return whether the activity analysis indicates that an equation is 'active'. :arg n: The :class:`int` index of the block. :arg i: The :class:`int` index of the equation. :returns: `True` if the equation is active, and `False` otherwise. .. !! processed by numpydoc !! .. py:method:: update(d, *, copy=True) Use the supplied :class:`Mapping` to update forward values. :arg d: A :class:`Mapping`. Updates values for those keys in `d` which are also in the :class:`.ReplayStorage`. :arg copy: Whether the values in `d` should be copied when being stored in the :class:`.ReplayStorage`. .. !! processed by numpydoc !! .. py:method:: popleft() Remove the first equation. Used to deallocate forward variables which are no longer needed as the solution of forward equations progresses. :returns: A :class:`tuple` `(n, i)`, indicating that equation `i` in block `n` has been removed. .. !! processed by numpydoc !! .. py:class:: Checkpoints Disk checkpointing abstract base class. .. !! processed by numpydoc !! .. py:method:: write(n, cp, data, storage) :abstractmethod: Write checkpoint data. :arg n: The :class:`int` index of the block with which the checkpoint data to be written is associated. :arg cp: See :meth:`.CheckpointStorage.update`. :arg data: See :meth:`.CheckpointStorage.update`. :arg storage: See :meth:`.CheckpointStorage.update`. .. !! processed by numpydoc !! .. py:method:: read(n, *, ics=True, data=True, ic_ids=None) :abstractmethod: Read checkpoint data. :arg n: The :class:`int` index of the block with which the checkpoint data to be read is associated. :arg ics: Whether forward restart data should be included. :arg data: Whether non-linear dependency data should be included. :arg ic_ids: A :class:`Container`. If provided then only variables with ID in `ic_ids` are included. :returns: A :class:`tuple` `(cp, data, storage)`. Elements of this :class:`tuple` are as for the three arguments for the :meth:`.CheckpointStorage.update` method. .. !! processed by numpydoc !! .. py:method:: delete(n) :abstractmethod: Delete checkpoint data. :arg n: The :class:`int` index of the block with which the checkpoint data to be deleted is associated. .. !! processed by numpydoc !! .. py:class:: PickleCheckpoints(prefix, *, comm=None) Disk checkpointing using the pickle module. :arg prefix: Checkpoint files are stored at `[prefix]_[root_pid]_[root_py2f]_[rank].pickle`. Here `prefix` is defined by this argument, `root_pid` is the process ID on the root process (i.e. process 0), `root_py2f` is the Fortran MPI communicator on the root process, and `rank` is the process rank. :arg comm: A communicator. .. !! processed by numpydoc !! .. py:method:: write(n, cp, data, storage) Write checkpoint data. :arg n: The :class:`int` index of the block with which the checkpoint data to be written is associated. :arg cp: See :meth:`.CheckpointStorage.update`. :arg data: See :meth:`.CheckpointStorage.update`. :arg storage: See :meth:`.CheckpointStorage.update`. .. !! processed by numpydoc !! .. py:method:: read(n, *, ics=True, data=True, ic_ids=None) Read checkpoint data. :arg n: The :class:`int` index of the block with which the checkpoint data to be read is associated. :arg ics: Whether forward restart data should be included. :arg data: Whether non-linear dependency data should be included. :arg ic_ids: A :class:`Container`. If provided then only variables with ID in `ic_ids` are included. :returns: A :class:`tuple` `(cp, data, storage)`. Elements of this :class:`tuple` are as for the three arguments for the :meth:`.CheckpointStorage.update` method. .. !! processed by numpydoc !! .. py:method:: delete(n) Delete checkpoint data. :arg n: The :class:`int` index of the block with which the checkpoint data to be deleted is associated. .. !! processed by numpydoc !! .. py:class:: HDF5Checkpoints(prefix, *, comm=None) Disk checkpointing using the h5py library. :arg prefix: Checkpoint files are stored at `[prefix]_[root_pid]_[root_py2f].hdf5`. Here `prefix` is defined by this argument, `root_pid` is the process ID on the root process (i.e. process 0), and `root_py2f` is the Fortran MPI communicator on the root process. :arg comm: A communicator. .. !! processed by numpydoc !! .. py:method:: write(n, cp, data, storage) Write checkpoint data. :arg n: The :class:`int` index of the block with which the checkpoint data to be written is associated. :arg cp: See :meth:`.CheckpointStorage.update`. :arg data: See :meth:`.CheckpointStorage.update`. :arg storage: See :meth:`.CheckpointStorage.update`. .. !! processed by numpydoc !! .. py:method:: read(n, *, ics=True, data=True, ic_ids=None) Read checkpoint data. :arg n: The :class:`int` index of the block with which the checkpoint data to be read is associated. :arg ics: Whether forward restart data should be included. :arg data: Whether non-linear dependency data should be included. :arg ic_ids: A :class:`Container`. If provided then only variables with ID in `ic_ids` are included. :returns: A :class:`tuple` `(cp, data, storage)`. Elements of this :class:`tuple` are as for the three arguments for the :meth:`.CheckpointStorage.update` method. .. !! processed by numpydoc !! .. py:method:: delete(n) Delete checkpoint data. :arg n: The :class:`int` index of the block with which the checkpoint data to be deleted is associated. .. !! processed by numpydoc !!