fdtdx.LinearReconstructEveryK

fdtdx.LinearReconstructEveryK#

class fdtdx.LinearReconstructEveryK(*, k=null, start_recording_after=0)[source]#

Bases: TimeStepFilter

Time step filter that performs linear reconstruction between sampled steps.

This filter saves field values every k time steps and uses linear interpolation to reconstruct values at intermediate time steps.

Quick Reference#

Attributes

Methods

Attributes#

LinearReconstructEveryK.k: int#

Number of time steps between saved values.

LinearReconstructEveryK.start_recording_after: int#

Time step to start recording from. Defaults to zero.

Methods#

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

LinearReconstructEveryK.compress(values, state, time_idx, key)[source]#

Compress field values at a given time step.

Parameters:

  • values (dict[str, jax.Array]) – Dictionary mapping field names to their values.

  • state (RecordingState) – Current recording state.

  • time_idx (jax.Array) – Current time step index.

  • key (jax.Array) – Random key for stochastic operations.

Returns:

Tuple containing:

  • Dictionary of compressed field values

  • Updated recording state

Return type:

tuple[dict[str, jax.Array], RecordingState]

LinearReconstructEveryK.decompress(values, state, arr_indices, time_idx, key)[source]#

Decompress field values to reconstruct data for a time step.

Parameters:

  • values (list[dict[str, jax.Array]]) – List of dictionaries containing array values needed for reconstruction.

  • state (RecordingState) – Current recording state.

  • arr_indices (jax.Array) – Array indices needed for reconstruction.

  • time_idx (jax.Array) – Time step index to reconstruct. scalar value.

  • key (jax.Array) – Random key for stochastic operations.

Returns:

Dictionary of reconstructed field values.

Return type:

dict[str, jax.Array]

LinearReconstructEveryK.get_class_fields()#
Return type:

list[TreeClassField]

LinearReconstructEveryK.get_public_fields()#
Return type:

list[TreeClassField]

LinearReconstructEveryK.indices_to_decompress(time_idx)[source]#

Get array indices needed to reconstruct data for a given time step.

Parameters:

time_idx (jax.Array) – Time step index to reconstruct.

Returns:

Array of indices needed to reconstruct the data for this time step.

Return type:

jax.Array

LinearReconstructEveryK.init_shapes(input_shape_dtypes, time_steps_max)[source]#

Initialize shapes and sizes for the time step filter.

Parameters:

  • input_shape_dtypes (dict[str, jax.ShapeDtypeStruct]) – Dictionary mapping field names to their shape/dtype information.

  • time_steps_max (int) – Maximum number of time steps in the simulation.

Returns:

A tuple containing:

  • Updated filter instance

  • Size of array for storing filtered data

  • Dictionary of data shapes/dtypes

  • Dictionary of state shapes/dtypes

Return type:

tuple[Self, int, dict[str, jax.ShapeDtypeStruct], dict[str, jax.ShapeDtypeStruct]]

LinearReconstructEveryK.time_to_array_index(time_idx)[source]#

Convert a time step index to its corresponding array index.

Parameters:

time_idx (int) – Time step index to convert.

Returns:

The corresponding array index if the time step is not filtered,

or -1 if the time step is filtered out.

Return type:

int

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