fdtdz documentation

Skip to main content
`Ctrl`+`K`

# fdtd-z API documentation

fdtdz_jax.fdtdz(epsilon, dt, source_field, source_waveform, source_position, absorption_mask, pml_kappa, pml_sigma, pml_alpha, pml_widths, output_steps, use_reduced_precision, launch_params)

  Execute a FDTD simulation.

  `fdtd-z` is an implementation of the finite-difference time-domain (FDTD) method that is efficiently mapped to the GPU via a systolic update scheme [1]. This function exposes, as a JAX primitive, a relatively low-level API to the underlying CUDA kernel.

  `fdtd-z` targets nanophotonic applications with a heavy prioritization on simulation throughput, as insufficient throughput is currently the bottleneck in many workflows (such as nanophotonic inverse design). As such, the flexibility and features of the engine have been kept to a bare minimum in the name of performance:

  - Reduced-precision mode is available and recommended which utilizes 16-bit floating point internally (input and output arrays are always at single-precision) and allows for ~2 times larger z-extent. Requires (Nvidia) GPUs of compute capability >= 6.0.

  - Z-extent (`zz`) of the simulation domain is fixed to a specific value to allow for a natural mapping to the 32 threads/warp architecture of Nvidia GPUs. The required value of `zz` is given by either `128 - p` or `64 - p` for the cases of reduced precision or full precision respectively, where `p` is the total number of PML layers for the simulation.

  - PML boundary conditions [2] are limited to the z-direction, adiabatic aborbing boundaries [3] must be used in the x- and y-directions.

  - Absorption (apart from absorbing boundary conditions) and dispersion are not implemented. The simulated structure must consist of a real-valued permittivity that depends only on the E-field component and spatial location.

  - The permeability is fixed at a value of `1` everywhere in the simulation domain.

  - Only current sources as hyperplanes along the y- and z-axes are implemented.

  `fdtd-z` uses a “dimensionless” system [4] where the permittivity and permeability of vacuum (and therefore the speed of light), as well as the size of the Yee cell are all set to a (dimensionless) value of `1` (although the size of the cell along the z-axis can be varied).

  [1]

  Lu, Jesse, and Jelena Vuckovic. “fdtd-z: A systolic scheme for GPU-accelerated nanophotonic simulation.” spinsphotonics/fdtdz

  [2]

  Roden, J. Alan, and Stephen D. Gedney. “Convolution PML (CPML): An efficient FDTD implementation of the CFS–PML for arbitrary media.” Microwave and optical technology letters 27.5 (2000): 334-339.

  [3]

  Oskooi, Ardavan, and Steven G. Johnson. “Distinguishing correct from incorrect PML proposals and a corrected unsplit PML for anisotropic, dispersive media.” Journal of Computational Physics 230.7 (2011): 2369-2377.

  [4]

  Oskooi, Ardavan F., et al. “MEEP: A flexible free-software package for electromagnetic simulations by the FDTD method.” Computer Physics Communications 181.3 (2010): 687-702.

  Parameters:

    - epsilon – `(3, xx, yy, zz)`-shaped array of floats representing the permittivity values at the `Ex`, `Ey`, and `Ez` nodes of the Yee cell respectively. We use the convention that these components are located at `(0.5, 0, 0)`, `(0, 0.5, 0)`, and `(0, 0, 0.5)` respectively, for a Yee cell of side-length 1.

    - dt – Scalar float representing the amount of time elapsed in one update step.

    - source_field – An array of shape `(2, 1, yy, zz)`, `(2, xx, 1, zz)`, or `(2, 2, xx, yy, 1)`-shaped array for a source at `x = source_position`, `y = source_position`, or `z = source_position` respectively. The `(2, 1, yy, zz)` source features `Ey` and `Ez` components in that order, while the `(2, xx, 1, zz)` source features `Ex` and `Ez` components in that order. The `(2, 2, xx, yy, 1)` allows for the specification of two separate source fields at `[0, :, :, :, :]` and `[1, :, :, :, :]` which each contain `Ex` and `Ey` components in that order.

    - source_waveform – `(tt, 2)`-shaped array of floats denoting the temporal variation to apply to the each of the source fields, where `tt` is the total number of update steps needed (note that the first update occurs at step `0`). Specifically, the subarray at `(tt, i)` applies a temporal variation to the `(i, 2, xx, yy. 1)` subarray of a source at `z = source_position`, while for a source at `x = source_position` or `y = source_position` the temporal variation is applied to the source field at `source_position - i`.

    - source_position – integer representing the position of the source along either the y- or z-axes. For the case of a source at `y = source_position` the source field is applied (with the corresponding waveform) at both `y = source_position` and `y = source_position - 1`, with the additional performance constraint that `source_position` must be even.

    - absorption_mask – `(3, xx, yy)`-shaped array of floats representing a z-invariant conductivity intended to allow for adiabatic absorbing boundary conditions along the x- and y-axes according to `conductivity(x, y, z) = aborption_mask(x, y) * epsilon(x, y, z)` for the `Ex`, `Ey`, and `Ez` component respectively.

    - pml_kappa – `(zz, 2)`-shaped array of floats denoting the distance along the z-axis between adjacent Yee cells. This is primarily intended to be used as a stretching parameter for the PML, but equivalently also determines the unit cell length along the z-axis throughout the simulation domain. Specifically, `(zz, 0)` represents the distance between successive layers of `Ex`, `Ey`, and `Hz` nodes, while `(zz, 1)` represents the distance between layers of `Hx`, `Hy`, and `Ez` nodes.

    - pml_sigma – `(zz, 2)`-shaped array of floats for the conductivity of the PML region, where `(zz, 0)` and `(zz, 1)` are the values at the (`Ex`, `Ey`, `Hz`) and (`Hx`, `Hy`, `Ez`) layers respectively. Must be set to `0` outside of the PML regions.

    - pml_alpha – `(zz, 2)`-shaped array of floats similar to `pml_sigma`. Must also be set to `0` outside of the PML regions.

    - pml_widths – `(bot, top)` integers specifying the number of cells which are to be designated as PML layers at the bottom and top of the simulation respectively. For performance reasons, the total number of PML layers used in the simulation (`bot + top`) is required to be a multiple of `4`.

    - output_steps – `(start, stop, interval)` tuple of integers denoting the update step at which to start recording output fields, the number of update steps separating successive output fields, and the step at which to stop recording (not included).

    - use_reduced_precision – If `True`, uses 16-bit (IEEE 754) precision for the simulation which allows for a maximum of 128 cells along the z-axis. Otherwise, uses 32-bit single-precision with a maximum of 64 cells along the z-axis. Both inputs and results are always expected as 32-bit arrays.

    - launch_params – Integers as an object in the form of `((blocku, blockv), (gridu, gridv), spacing, (cc_major, cc_minor))`, specifying the structure of the systolic update to use on the GPU where `(blocku, blockv)` determines the layout of warps in the u- and v-directions within a block and should be `(2, 4)` or `(4, 2)`; `(gridu, gridv)` specify the layout of blocks on the GPU and must be equal to or less than the number of streaming multiprocessors on the GPU because of the need for grid-wide synchronization; `spacing` controls the number of buffers used between each block and its downstream neighbor and should be tuned to balance between reducing grid synchronization overhead and staying within the limits of the L2 cache; and `(cc_major, cc_minor)` major and minor compute capability of the device. Used to determine which precompiled kernel to use. Currently allowed values are `(3, 7)`, `(6, 0)`, `(7, 0)`, `(7, 5)`, and `(8, 0)`. Recommended to use the latest compute capability kernel possible that does not exceed the compute capability of the device.

  Returns:

    `(n, 3, xx, yy, zz)` array of floats representing `n` output fields, where each output field consists of the values of the `Ex`, `Ey`, and `Ez` node (in that order) over the simulation domain, at a specific update step.