tlm_adjoint.checkpointing

Module Contents

class tlm_adjoint.checkpointing.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 int indicating the block index and i is an int indicating the equation index within that block. Non-linear-dependency data for an Equation can be accessed via, e.g.

nl_deps = cp[(n, i)]

where cp is a CheckpointStorage. Here nl_deps is a tuple of variables storing values associated with eq.nonlinear_dependencies(), for Equation i in block n.

Parameters:

• store_ics – Whether to enable storage of forward restart data.

• store_data – Whether to enable storage of non-linear dependency data.

property store_ics

Whether storage of forward restart data is enabled.

property store_data

Whether storage of non-linear dependency data is enabled.

configure(*, store_ics, store_data)

Enable or disable storage of forward restart and non-linear dependency data.

Parameters:

• store_ics – Whether storage of forward restart data should be enabled (store_ics=True) or disabled (store_ics=False).

• store_data – Whether storage of non-linear dependency data should be enabled (store_data=True) or disabled (store_data=False).

clear(*, clear_ics=True, clear_data=True, clear_refs=False)

Clear stored data.

Parameters:

• clear_ics – Whether forward restart data should be cleared.

• clear_data – Whether non-linear dependency data should be cleared.

• clear_refs – Whether references should be cleared.

initial_condition(x, *, copy=True)

Return the forward restart value associated with a variable x.

Parameters:

• x – The variable, or the int variable ID, for which the forward restart value should be returned.

• 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.

initial_conditions(*, cp=True, refs=False, copy=True)

Access stored forward restart data.

Parameters:

• cp – Whether to include forward restart data that is stored by value.

• refs – Whether to include data that is stored by reference. May include non-linear dependency data that is not forward restart data.

• copy – If True then a copy of the stored data is returned. If False then internal variables storing the data are returned.

Returns:

A VariableStateLockDictionary, with items (x_id: x_value), where x_id is the int variable ID and x_value is a variable storing the data.

add_initial_condition(x, value=None)

Store forward restart data / an initial condition dependency.

Parameters:

• x – The initial condition dependency variable.

• value – A variable defining the initial condition dependency value. x is used if not supplied.

update_keys(n, i, eq)

The 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.

Parameters:

• n – The int index of the block.

• i – The int index of the equation.

• eq – An Equation, equation i in block n.

add_equation(n, i, eq, *, deps=None, nl_deps=None)

Store checkpoint data associated with an equation.

Parameters:

• n – The int index of the block.

• i – The int index of the equation.

• eq – An Equation, equation i in block n.

• deps – Equation dependency values. eq.dependencies() is used if not supplied.

• nl_deps – Equation non-linear dependency values. Extracted from deps if not supplied.

add_equation_data(n, i, eq, *, nl_deps=None)

Store checkpoint data associated with an equation. As CheckpointStorage.add_equation(), but adds only non-linear dependency data.

Parameters:

• n – The int index of the block.

• i – The int index of the equation.

• eq – An Equation, equation i in block n.

• nl_deps – Equation non-linear dependency values. eq.nonlinear_dependencies() is used if not supplied.

checkpoint_data(*, ics=True, data=True, copy=True)

Extract checkpoint data.

Parameters:

• ics – Whether to extract forward restart data.

• data – Whether to extract non-linear dependency data.

• copy – If True then a copy of the stored data is returned. If False then internal variables storing the data are returned.

Returns:

A tuple (cp, data, storage). Elements of this tuple are as for the three arguments for the CheckpointStorage.update() method, and here storage is a VariableStateLockDictionary.

update(cp, data, storage, *, copy=True)

Update the CheckpointStorage using the provided checkpoint data. Used to update the CheckpointStorage from loaded data. Note that the CheckpointStorage is not cleared prior to updating using the provided data.

Parameters:

• cp – A tuple of keys. Forward restart data is defined by (storage[key] for key in cp).

• data – A 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).

• storage – The stored data. A Mapping with items ((x_id, x_indices), x_value). x_id is the 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.

• copy – Whether the values in storage should be copied when being stored in the CheckpointStorage.

class tlm_adjoint.checkpointing.ReplayStorage(blocks, N0, N1, *, transpose_deps=None)

Storage used when solving forward equations.

A value for a forward variable can be accessed via

x_value = replay_storage[x]

and set via

replay_storage[x] = x_value

where here x is either a variable of an int variable ID. Containment can also be tested,

if x in replay_storage:
    [...]

Parameters:

• blocks – A Sequence or Mapping, whose elements or values are Sequence objects containing Equation objects. Forward equations.

• N0 – An int. (blocks[n] for n in range(N0, N1)) defines the forward equations which will be solved.

• N1 – An int. (blocks[n] for n in range(N0, N1)) defines the forward equations which will be solved.

• transpose_deps – A TransposeComputationalGraph. If supplied then an activity analysis is applied.

is_active(n, i)

Return whether the activity analysis indicates that an equation is ‘active’.

Parameters:

• n – The int index of the block.

• i – The int index of the equation.

Returns:

True if the equation is active, and False otherwise.

update(d, *, copy=True)

Use the supplied Mapping to update forward values.

Parameters:

• d – A Mapping. Updates values for those keys in d which are also in the ReplayStorage.

• copy – Whether the values in d should be copied when being stored in the ReplayStorage.

popleft()

Remove the first equation. Used to deallocate forward variables which are no longer needed as the solution of forward equations progresses.

Returns:

A tuple (n, i), indicating that equation i in block n has been removed.

class tlm_adjoint.checkpointing.Checkpoints

Disk checkpointing abstract base class.

abstract write(n, cp, data, storage)

Write checkpoint data.

Parameters:

• n – The int index of the block with which the checkpoint data to be written is associated.

• cp – See CheckpointStorage.update().

• data – See CheckpointStorage.update().

• storage – See CheckpointStorage.update().

abstract read(n, *, ics=True, data=True, ic_ids=None)

Read checkpoint data.

Parameters:

• n – The int index of the block with which the checkpoint data to be read is associated.

• ics – Whether forward restart data should be included.

• data – Whether non-linear dependency data should be included.

• ic_ids – A Container. If provided then only variables with ID in ic_ids are included.

Returns:

A tuple (cp, data, storage). Elements of this tuple are as for the three arguments for the CheckpointStorage.update() method.

abstract delete(n)

Delete checkpoint data.

Parameters:

n – The int index of the block with which the checkpoint data to be deleted is associated.

class tlm_adjoint.checkpointing.PickleCheckpoints(prefix, *, comm=None)

Disk checkpointing using the pickle module.

Parameters:

• 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.

• comm – A communicator.

write(n, cp, data, storage)

Write checkpoint data.

Parameters:

• n – The int index of the block with which the checkpoint data to be written is associated.

• cp – See CheckpointStorage.update().

• data – See CheckpointStorage.update().

• storage – See CheckpointStorage.update().

read(n, *, ics=True, data=True, ic_ids=None)

Read checkpoint data.

Parameters:

• n – The int index of the block with which the checkpoint data to be read is associated.

• ics – Whether forward restart data should be included.

• data – Whether non-linear dependency data should be included.

• ic_ids – A Container. If provided then only variables with ID in ic_ids are included.

Returns:

A tuple (cp, data, storage). Elements of this tuple are as for the three arguments for the CheckpointStorage.update() method.

delete(n)

Delete checkpoint data.

Parameters:

n – The int index of the block with which the checkpoint data to be deleted is associated.

class tlm_adjoint.checkpointing.HDF5Checkpoints(prefix, *, comm=None)

Disk checkpointing using the h5py library.

Parameters:

• 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.

• comm – A communicator.

write(n, cp, data, storage)

Write checkpoint data.

Parameters:

• n – The int index of the block with which the checkpoint data to be written is associated.

• cp – See CheckpointStorage.update().

• data – See CheckpointStorage.update().

• storage – See CheckpointStorage.update().

read(n, *, ics=True, data=True, ic_ids=None)

Read checkpoint data.

Parameters:

• n – The int index of the block with which the checkpoint data to be read is associated.

• ics – Whether forward restart data should be included.

• data – Whether non-linear dependency data should be included.

• ic_ids – A Container. If provided then only variables with ID in ic_ids are included.

Returns:

A tuple (cp, data, storage). Elements of this tuple are as for the three arguments for the CheckpointStorage.update() method.

delete(n)

Delete checkpoint data.

Parameters:

n – The int index of the block with which the checkpoint data to be deleted is associated.