fdtdx.GradientConfig

class fdtdx.GradientConfig(*, method='reversible', recorder=None, num_checkpoints=None, num_checkpoints_reversible=0)

Bases: TreeClass

Configuration for gradient computation in simulations.

This class handles settings for automatic differentiation, supporting either invertible differentiation with a recorder or checkpointing-based differentiation.

Quick Reference

Attributes

• method

• num_checkpoints

• num_checkpoints_reversible

• recorder

Methods

• aset

• get_class_fields

• get_public_fields

Attributes

GradientConfig.method: Literal['reversible', 'checkpointed']

Method for gradient computation. Can be either “reversible” when using the time reversible autodiff, or “checkpointed” for the exact checkpointing algorithm.

GradientConfig.num_checkpoints: int | None

Optional number of checkpoints for checkpointing-based differentiation. Needs to be provided for checkpointing gradient computation. Defaults to None.

GradientConfig.num_checkpoints_reversible: int

Number of interior full-field checkpoints for the "reversible" method. The reversible backward pass reconstructs the field state by running the simulation in reverse; for lossy/dispersive materials this reverse reconstruction can accumulate numerical error over the full trajectory. Setting this to k - 1 partitions the run into k slices and stores a full-field checkpoint at each interior slice boundary during the forward pass. The backward pass then resets the reverse reconstruction to the exact checkpoint at every boundary, bounding the reconstruction drift to a single slice (~time_steps_total / k steps) at the cost of O(k) field memory. The default 0 reproduces the classic single full reverse pass (no interior checkpoints; only the final field, which is available for free, is used). Ignored by the "checkpointed" method. Must not exceed time_steps_total - 1.

GradientConfig.recorder: Recorder | None

Optional recorder for invertible differentiation. Needs to be provided for reversible autodiff. Defaults to None

Methods

GradientConfig.aset(attr_name, val, create_new_ok=False)

Sets an attribute of this class. In contrast to the classical .at[].set(), this method updates the class attribute directly and does not only operate on jax pytree leaf nodes. Instead, replaces the full attribute with the new value.

The attribute can either be the attribute name of this class, or for nested classes it can also be the attribute name of a class, which itself is an attribute of this class. The syntax for this operation could look like this: “a->b->[0]->[‘name’]”. Here, the current class has an attribute a, which has an attribute b, which is a list, which we index at index 0, which is an element of type dictionary, which we index using the dictionary key ‘name’.

Note that dictionary keys cannot contain square brackets or single quotes (even if they are escaped).

Parameters:

• attr_name (str) – Name of attribute to set

• val (Any) – Value to set the attribute to

• create_new_ok (bool, optional) – If false (default), throw an error if the attribute does not exist. If true, creates a new attribute if the attribute name does not exist yet.

Returns:

Updated instance with new attribute value

Return type:

Self

GradientConfig.get_class_fields()

Return type:

list[TreeClassField]

GradientConfig.get_public_fields()

Return type:

list[TreeClassField]

If you find any errors in the documentation, please report them in the Github Issues!