Skip to content

Configuration

fdtdx.core.config.GradientConfig

Bases: ExtendedTreeClass

Configuration for gradient computation in simulations.

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

Parameters:

Name Type Description Default
recorder Recorder | None

Optional recorder for invertible differentiation. If provided, invertible differentiation will be used.

 None
num_checkpoints int | None

Optional number of checkpoints for checkpointing-based differentiation. If provided, checkpointing will be used.

 None

Raises:

Type Description
Exception

If both recorder and num_checkpoints are provided, or if neither is provided.
Source code in src/fdtdx/core/config.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
class GradientConfig(ExtendedTreeClass):
    """Configuration for gradient computation in simulations.

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

    Args:
        recorder: Optional recorder for invertible differentiation. If provided,
            invertible differentiation will be used.
        num_checkpoints: Optional number of checkpoints for checkpointing-based
            differentiation. If provided, checkpointing will be used.

    Raises:
        Exception: If both recorder and num_checkpoints are provided, or if neither
            is provided.
    """

    def __init__(
        self,
        recorder: Recorder | None = None,  # if not none, use invertible diff
        num_checkpoints: int | None = None,
    ):
        self.recorder = recorder
        self.num_checkpoints = num_checkpoints
        if self.recorder is not None and self.num_checkpoints is not None:
            raise Exception("Cannot use both invertible and checkpointing autodiff!")
        if self.recorder is None and self.num_checkpoints is None:
            raise Exception("Need either recorder or checkpoints to define autograd!")

fdtdx.core.config.SimulationConfig

Bases: ExtendedTreeClass

Configuration settings for FDTD simulations.

This class contains all the parameters needed to configure and run an FDTD simulation, including spatial and temporal discretization, hardware backend, and gradient computation settings.

Attributes:

Name Type Description
time float

Total simulation time in seconds.
resolution float

Spatial resolution of the simulation grid in meters.
backend Literal['gpu', 'tpu', 'cpu']

Computation backend ('gpu', 'tpu', or 'cpu').
dtype dtype

Data type for numerical computations.
courant_factor float

Safety factor for the Courant condition (default: 0.99).
gradient_config GradientConfig | None

Optional configuration for gradient computation.
Source code in src/fdtdx/core/config.py 
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
@extended_autoinit
class SimulationConfig(ExtendedTreeClass):
    """Configuration settings for FDTD simulations.

    This class contains all the parameters needed to configure and run an FDTD
    simulation, including spatial and temporal discretization, hardware backend,
    and gradient computation settings.

    Attributes:
        time: Total simulation time in seconds.
        resolution: Spatial resolution of the simulation grid in meters.
        backend: Computation backend ('gpu', 'tpu', or 'cpu').
        dtype: Data type for numerical computations.
        courant_factor: Safety factor for the Courant condition (default: 0.99).
        gradient_config: Optional configuration for gradient computation.
    """

    time: float
    resolution: float
    backend: Literal["gpu", "tpu", "cpu"] = frozen_field(default="gpu")
    dtype: jnp.dtype = frozen_field(default=jnp.float32)
    courant_factor: float = 0.99
    gradient_config: GradientConfig | None = None

    def __post_init__(self):
        if self.backend in ["gpu", "tpu"]:
            # Try to initialize GPU
            try:
                jax.devices(self.backend)
                logger.info(f"{str.upper(self.backend)} found and will be used for computations")
                jax.config.update("jax_platform_name", self.backend)
            except RuntimeError:
                logger.warning(f"{str.upper(self.backend)} not found, falling back to CPU!")
                self.backend = "cpu"
                jax.config.update("jax_platform_name", "cpu")
        else:
            jax.config.update("jax_platform_name", "cpu")

    @property
    def courant_number(self) -> float:
        """Calculate the Courant number for the simulation.

        The Courant number is a dimensionless quantity that determines stability
        of the FDTD simulation. It represents the ratio of the physical propagation
        speed to the numerical propagation speed.

        Returns:
            float: The Courant number, scaled by the courant_factor and normalized
                for 3D simulations.
        """
        return self.courant_factor / math.sqrt(3)

    @property
    def time_step_duration(self) -> float:
        """Calculate the duration of a single time step.

        The time step duration is determined by the Courant condition to ensure
        numerical stability. It depends on the spatial resolution and the speed
        of light.

        Returns:
            float: Time step duration in seconds, calculated using the Courant
                condition and spatial resolution.
        """
        return self.courant_number * self.resolution / constants.c

    @property
    def time_steps_total(self) -> int:
        """Calculate the total number of time steps for the simulation.

        Determines how many discrete time steps are needed to simulate the
        specified total simulation time, based on the time step duration.

        Returns:
            int: Total number of time steps needed to reach the specified
                simulation time.
        """
        return round(self.time / self.time_step_duration)

    @property
    def max_travel_distance(self) -> float:
        """Calculate the maximum distance light can travel during the simulation.

        This represents the theoretical maximum distance that light could travel
        through the simulation volume, useful for determining if the simulation
        time is sufficient for light to traverse the entire domain.

        Returns:
            float: Maximum travel distance in meters, based on the speed of light
                and total simulation time.
        """
        return constants.c * self.time

    @property
    def only_forward(self) -> bool:
        """Check if the simulation is forward-only (no gradient computation).

        Forward-only simulations don't compute gradients and are used when only
        the forward propagation of electromagnetic fields is needed, without
        optimization.

        Returns:
            bool: True if no gradient configuration is specified, False otherwise.
        """
        return self.gradient_config is None

    @property
    def invertible_optimization(self) -> bool:
        """Check if invertible optimization is enabled.

        Invertible optimization uses time-reversibility of Maxwell's equations
        to compute gradients with reduced memory requirements compared to
        checkpointing-based methods.

        Returns:
            bool: True if gradient computation uses invertible differentiation
                (recorder is specified), False otherwise.
        """
        if self.gradient_config is None:
            return False
        return self.gradient_config.recorder is not None

courant_number: float property

Calculate the Courant number for the simulation.

The Courant number is a dimensionless quantity that determines stability of the FDTD simulation. It represents the ratio of the physical propagation speed to the numerical propagation speed.

Returns:

Name Type Description
float float

The Courant number, scaled by the courant_factor and normalized for 3D simulations.

invertible_optimization: bool property

Check if invertible optimization is enabled.

Invertible optimization uses time-reversibility of Maxwell's equations to compute gradients with reduced memory requirements compared to checkpointing-based methods.

Returns:

Name Type Description
bool bool

True if gradient computation uses invertible differentiation (recorder is specified), False otherwise.

max_travel_distance: float property

Calculate the maximum distance light can travel during the simulation.

This represents the theoretical maximum distance that light could travel through the simulation volume, useful for determining if the simulation time is sufficient for light to traverse the entire domain.

Returns:

Name Type Description
float float

Maximum travel distance in meters, based on the speed of light and total simulation time.

only_forward: bool property

Check if the simulation is forward-only (no gradient computation).

Forward-only simulations don't compute gradients and are used when only the forward propagation of electromagnetic fields is needed, without optimization.

Returns:

Name Type Description
bool bool

True if no gradient configuration is specified, False otherwise.

time_step_duration: float property

Calculate the duration of a single time step.

The time step duration is determined by the Courant condition to ensure numerical stability. It depends on the spatial resolution and the speed of light.

Returns:

Name Type Description
float float

Time step duration in seconds, calculated using the Courant condition and spatial resolution.

time_steps_total: int property

Calculate the total number of time steps for the simulation.

Determines how many discrete time steps are needed to simulate the specified total simulation time, based on the time step duration.

Returns:

Name Type Description
int int

Total number of time steps needed to reach the specified simulation time.