Detector and Interface Management

Detector Updates

fdtdx.fdtd.update.update_detector_states(time_step, arrays, objects, H_prev, inverse)

Updates detector states based on current field values.

Handles field interpolation for accurate detector measurements. By default, interpolation is disabled for performance during optimization, but can be enabled for final evaluation. Interpolation is needed due to the staggered nature of E and H fields on the Yee grid.

Parameters:

Name	Type	Description	Default
time_step	Array	Current simulation time step	required
arrays	ArrayContainer	Container with E, H fields and material properties	required
objects	ObjectContainer	Container with detectors and other simulation objects	required
H_prev	Array	Previous H field values for interpolation	required
inverse	bool	Whether this is a forward or reverse update	required

Returns:

Type	Description
ArrayContainer	Updated ArrayContainer with new detector states

Source code in src/fdtdx/fdtd/update.py

def update_detector_states(
    time_step: jax.Array,
    arrays: ArrayContainer,
    objects: ObjectContainer,
    H_prev: jax.Array,
    inverse: bool,
) -> ArrayContainer:
    """Updates detector states based on current field values.

    Handles field interpolation for accurate detector measurements. By default,
    interpolation is disabled for performance during optimization, but can be
    enabled for final evaluation. Interpolation is needed due to the staggered
    nature of E and H fields on the Yee grid.

    Args:
        time_step: Current simulation time step
        arrays: Container with E, H fields and material properties
        objects: Container with detectors and other simulation objects
        H_prev: Previous H field values for interpolation
        inverse: Whether this is a forward or reverse update

    Returns:
        Updated ArrayContainer with new detector states
    """
    """Updates detector states based on current field values.

    Handles field interpolation for accurate detector measurements. By default,
    interpolation is disabled for performance during optimization, but can be
    enabled for final evaluation. Interpolation is needed due to the staggered
    nature of E and H fields on the Yee grid.

    Args:
        time_step: Current simulation time step
        arrays: Container with E, H fields and material properties
        objects: Container with detectors and other simulation objects
        H_prev: Previous H field values for interpolation
        inverse: Whether this is a forward or reverse update

    Returns:
        Updated ArrayContainer with new detector states
    """
    interpolated_E, interpolated_H = interpolate_fields(
        E_field=arrays.E,
        H_field=(H_prev + arrays.H) / 2,
    )

    def helper_fn(E_input, H_input, detector: Detector):
        detector = tc.unfreeze(detector)
        return detector.update(
            time_step=time_step,
            E=E_input,
            H=H_input,
            state=arrays.detector_states[detector.name],
            inv_permittivity=arrays.inv_permittivities,
            inv_permeability=arrays.inv_permeabilities,
        )

    state = arrays.detector_states
    to_update = objects.backward_detectors if inverse else objects.forward_detectors
    for d in to_update:
        state[d.name] = jax.lax.cond(
            d._is_on_at_time_step_arr[time_step],
            helper_fn,
            lambda e, h, _: state[d.name],
            interpolated_E if d.exact_interpolation else arrays.E,
            interpolated_H if d.exact_interpolation else arrays.H,
            tc.freeze(d),
        )
    arrays = arrays.aset("detector_states", state)
    return arrays

Interface Handling

fdtdx.fdtd.update.collect_interfaces(time_step, arrays, objects, config, key)

Collects field values at PML interfaces for gradient computation.

Part of the memory-efficient automatic differentiation implementation. Saves field values at boundaries between PML and inner simulation volume since PML updates are not time-reversible.

Parameters:

Name	Type	Description	Default
time_step	Array	Current simulation time step	required
arrays	ArrayContainer	Container with fields and material properties	required
objects	ObjectContainer	Container with PML and other simulation objects	required
config	SimulationConfig	Simulation configuration with gradient settings	required
key	Array	Random key for compression	required

Returns:

Type	Description
ArrayContainer	Updated ArrayContainer with recorded interface values

Source code in src/fdtdx/fdtd/update.py

def collect_interfaces(
    time_step: jax.Array,
    arrays: ArrayContainer,
    objects: ObjectContainer,
    config: SimulationConfig,
    key: jax.Array,
) -> ArrayContainer:
    """Collects field values at PML interfaces for gradient computation.

    Part of the memory-efficient automatic differentiation implementation.
    Saves field values at boundaries between PML and inner simulation volume
    since PML updates are not time-reversible.

    Args:
        time_step: Current simulation time step
        arrays: Container with fields and material properties
        objects: Container with PML and other simulation objects
        config: Simulation configuration with gradient settings
        key: Random key for compression

    Returns:
        Updated ArrayContainer with recorded interface values
    """
    """Collects field values at PML interfaces for gradient computation.

    Part of the memory-efficient automatic differentiation implementation.
    Saves field values at boundaries between PML and inner simulation volume
    since PML updates are not time-reversible.

    Args:
        time_step: Current simulation time step
        arrays: Container with fields and material properties
        objects: Container with PML and other simulation objects
        config: Simulation configuration with gradient settings
        key: Random key for compression

    Returns:
        Updated ArrayContainer with recorded interface values
    """
    if config.gradient_config is None or config.gradient_config.recorder is None:
        raise Exception("Need recorder to record boundaries")
    if arrays.recording_state is None:
        raise Exception("Need recording state to record boundaries")
    values = collect_boundary_interfaces(
        arrays=arrays,
        pml_objects=objects.pml_objects,
    )
    recording_state = config.gradient_config.recorder.compress(
        values=values,
        state=arrays.recording_state,
        time_step=time_step,
        key=key,
    )
    arrays = arrays.aset("recording_state", recording_state)
    return arrays

fdtdx.fdtd.update.add_interfaces(time_step, arrays, objects, config, key)

Adds previously collected interface values back to the fields.

Part of the memory-efficient automatic differentiation implementation. Restores saved field values at PML boundaries during reverse propagation since PML updates are not time-reversible.

Parameters:

Name	Type	Description	Default
time_step	Array	Current simulation time step	required
arrays	ArrayContainer	Container with fields and material properties	required
objects	ObjectContainer	Container with PML and other simulation objects	required
config	SimulationConfig	Simulation configuration with gradient settings	required
key	Array	Random key for decompression	required

Returns:

Type	Description
ArrayContainer	Updated ArrayContainer with restored interface values

Source code in src/fdtdx/fdtd/update.py

def add_interfaces(
    time_step: jax.Array,
    arrays: ArrayContainer,
    objects: ObjectContainer,
    config: SimulationConfig,
    key: jax.Array,
) -> ArrayContainer:
    """Adds previously collected interface values back to the fields.

    Part of the memory-efficient automatic differentiation implementation.
    Restores saved field values at PML boundaries during reverse propagation
    since PML updates are not time-reversible.

    Args:
        time_step: Current simulation time step
        arrays: Container with fields and material properties
        objects: Container with PML and other simulation objects
        config: Simulation configuration with gradient settings
        key: Random key for decompression

    Returns:
        Updated ArrayContainer with restored interface values
    """
    """Adds previously collected interface values back to the fields.

    Part of the memory-efficient automatic differentiation implementation.
    Restores saved field values at PML boundaries during reverse propagation
    since PML updates are not time-reversible.

    Args:
        time_step: Current simulation time step
        arrays: Container with fields and material properties
        objects: Container with PML and other simulation objects
        config: Simulation configuration with gradient settings
        key: Random key for decompression

    Returns:
        Updated ArrayContainer with restored interface values
    """
    if config.gradient_config is None or config.gradient_config.recorder is None:
        raise Exception("Need recorder to record boundaries")
    if arrays.recording_state is None:
        raise Exception("Need recording state to record boundaries")

    values, state = config.gradient_config.recorder.decompress(
        state=arrays.recording_state,
        time_step=time_step,
        key=key,
    )
    arrays = arrays.aset("recording_state", state)

    container = add_boundary_interfaces(
        arrays=arrays,
        values=values,
        pml_objects=objects.pml_objects,
    )

    return container